|
| 1 | +// https://atcoder.jp/contests/agc026/tasks/agc026_c |
| 2 | +// |
| 3 | +// 以下のクレートを使用。 |
| 4 | +// |
| 5 | +// - `either` |
| 6 | +// - `itertools` |
| 7 | +// - `proconio` |
| 8 | +// - `rustc-hash` |
| 9 | + |
| 10 | +use either::Either; |
| 11 | +use itertools::Itertools as _; |
| 12 | +use proconio::input; |
| 13 | +use proconio::marker::Bytes; |
| 14 | +use rustc_hash::FxHashMap; |
| 15 | +use std::collections::HashMap; |
| 16 | + |
| 17 | +fn main() { |
| 18 | + // `proconio::input!`。 |
| 19 | + // |
| 20 | + // https://docs.rs/proconio/0.3.6/proconio/macro.input.html |
| 21 | + input! { |
| 22 | + n: usize, |
| 23 | + s: Bytes, |
| 24 | + } |
| 25 | + |
| 26 | + let first = pairs(&s[..n]); |
| 27 | + |
| 28 | + // `itertools::Itertools`により`.collect::<Vec<_>>()`を`.collect_vec`と書ける。 |
| 29 | + // 結構微妙だけど`Itertools`を`use`してるなら考えてもいいかもしれない。 |
| 30 | + // |
| 31 | + // https://docs.rs/itertools/0.8/itertools/trait.Itertools.html#method.collect_vec |
| 32 | + let second = pairs(&s[n..].iter().copied().rev().collect_vec()); |
| 33 | + |
| 34 | + let ans = first |
| 35 | + .into_iter() |
| 36 | + .flat_map(|(p, n1)| second.get(&p).map(|n2| n1 * n2)) |
| 37 | + .sum::<u64>(); |
| 38 | + println!("{}", ans); |
| 39 | +} |
| 40 | + |
| 41 | +// `rustc-hash`が使えるのだが、この問題においてはデフォルトのハッシャーとの違いがよくわからない。 |
| 42 | +// というよりハッシュがボトルネックではない気がする。 |
| 43 | +// |
| 44 | +// https://docs.rs/rustc-hash/1 |
| 45 | +fn pairs(half: &[u8]) -> FxHashMap<(Vec<u8>, Vec<u8>), u64> { |
| 46 | + let mut counter = HashMap::default(); |
| 47 | + for bits in 0..1 << half.len() { |
| 48 | + // `<_ as itertools::Itertools>::partition_map`を用いることで2^{`half`}を表わすビット列に対して、 |
| 49 | + // 具体的な(赤の列, 青の列) / (青の列, 赤の列)を求める。 |
| 50 | + // |
| 51 | + // https://docs.rs/itertools/0.8/itertools/trait.Itertools.html#method.partition_map |
| 52 | + let pair = half.iter().enumerate().partition_map(|(i, &c)| { |
| 53 | + if bits >> i & 1 == 1 { |
| 54 | + Either::Left(c) |
| 55 | + } else { |
| 56 | + Either::Right(c) |
| 57 | + } |
| 58 | + }); |
| 59 | + *counter.entry(pair).or_insert(0) += 1; |
| 60 | + } |
| 61 | + counter |
| 62 | +} |
0 commit comments