AtCoder AGC 035 A - XOR Circle ※正攻法 (Rust)

irisruneです。一応Rust第2回です。

問題

atcoder.jp

300点問題ですが、正攻法で解こうとすると考慮すべきケースを整理したり実装がやや面倒だったりと厄介な問題です。

use std::collections::HashMap;

fn main() {
    input!{
        n: usize,
        la: [usize; n]
    }
    let mut map: HashMap<usize, usize> = HashMap::new();
    for a in la {
        let value = map.entry(a).or_insert(0);
        *value += 1;
    }

    if map.len() >= 4 {
        println!("No");
        return;
    }

    let mut keys = vec![];
    let mut values = vec![];
    for (key, val) in map.iter() {
        keys.push(*key);
        values.push(*val);
    }

    if map.len() == 1 {
        if keys[0] == 0 {
            println!("Yes");
        }else {
            println!("No");
        }
        return;
    }

    if map.len() == 2 {
        if keys[0] == 0 && values[0] * 2 == values[1] {
            println!("Yes");
            return;
        } else if keys[1] == 0 && values[1] * 2 == values[0] {
            println!("Yes");
            return;
        }
        println!("No");
        return;
    }

    if values[0] != values[1] || values[1] != values[2] {
        println!("No");
        return;
    }
    if keys[0] ^ keys[1] == keys[2] {
        println!("Yes");
        return;
    }
    println!("No");
}

まずこの問題の最も大きなポイントとしては、条件を満たす場合に登場する帽子の値は高々3種類までということです。これは排他的論理和の性質を用いることで導き出すことができます。

問題はある値の帽子が何個存在する必要があるかを考える部分で、登場する帽子の値の種類数によって変わってきます。具体的には、

• 3種類であればそれぞれ同数（）であればよいことが簡単に導き出せます。
• 2種類の場合は、1種類は値が0である必要があり、さらにもう1種類の半分の個数（）存在する必要があります。
• 1種類の場合は、個数はでよいですが値は0である必要があります。

以上のように場合分けを行う必要があり、かなり見落としやすく感じられました。実際、当初は3種類の場合しか正しく考えられておらず、2回WAを出しています。

また、実装についても少し面倒なところがあり、上記コードのように連想配列を使う方法では連想配列で帽子の値と個数を整理した後に、一度配列（ベクトル）に移してから条件を満たすか確認する必要があります。

公式解説では計算量と書かれているので、おそらくソートを行ってから順番に確認する方法を取るのだと思いますが、いずれにしてもやや面倒かつ冗長な気がしています。

雑記

Rustについて、

• Hashmap.entry().or_insert()で条件分岐の必要なく連想配列への値の追加が行えるのは便利だと思いました。
• 一方で参照渡しと値渡しを的確に使い分けないとコンパイルすら許されないなど、プログラムを動かすまでの道のりが険しく感じられます。
  • プログラムが正しく動くことがある程度保証されるのは強みだと思いますが、個人的には動かしてからデバッグしたいので少し合わない所があるかもしれません。