AtCoder Beginner Contest 384 F

https://atcoder.jp/contests/abc384/tasks/abc384_f

ふつうに計算すれば以上かかりますが、Aを偶数と奇数に分けてみましょう。そうすると、偶数から一つ奇数から一つ取って足せば割らなくてもいいので、和を簡単に計算できます。これがd=2のときです。d=4のときは、4で割った余りで分類して、余りを足して2または6になるペアで簡単に和が取れます。これをd=8、16と処理していき、dになるペアがなくなったら終わりです。

// Double Sum 2
#![allow(non_snake_case)]

use std::collections::HashMap;


//////////////////// library ////////////////////

fn read<T: std::str::FromStr>() -> T {
    let mut line = String::new();
    std::io::stdin().read_line(&mut line).ok();
    line.trim().parse().ok().unwrap()
}

fn read_vec<T: std::str::FromStr>() -> Vec<T> {
    read::<String>().split_whitespace()
            .map(|e| e.parse().ok().unwrap()).collect()
}


//////////////////// process ////////////////////

fn read_input() -> Vec<i64> {
    let _N: usize = read();
    let A: Vec<i64> = read_vec();
    A
}

fn div_2pow(mut n: i64) -> i64 {
    while n % 2 == 0 {
        n /= 2
    }
    n
}

fn sum_same(v: &Vec<i64>) -> i64 {
    v.iter().map(|&e| e).sum::<i64>() * (v.len() as i64 - 1)
}

fn sum_diff(v: &Vec<i64>, w: &Vec<i64>) -> i64 {
    v.iter().map(|&e| e).sum::<i64>() * (w.len() as i64) +
    w.iter().map(|&e| e).sum::<i64>() * (v.len() as i64)
}

fn classify(v: &Vec<i64>, d: i64) -> HashMap<i64, Vec<i64>> {
    let mut dic: HashMap<i64, Vec<i64>> = HashMap::new();
    for &n in v.iter() {
        let r = n & (d - 1);
        let e = dic.entry(r).or_insert(vec![]);
        (*e).push(n)
    }
    dic
}

fn is_remained(dic: &HashMap<i64, Vec<i64>>, d: i64) -> bool {
    for (&r1, ns1) in dic.into_iter() {
        let r2 = if r1 == 0 { 0 } else { d - r1 };
        if r1 == r2 {
            if ns1.len() > 1 {
                return true
            }
        }
        else if r1 < r2 {
            if dic.contains_key(&r2) {
                return true
            }
        }
    }
    false
}

fn F(A: Vec<i64>) -> i64 {
    let mut s = A.iter().map(|&n| div_2pow(n)).sum::<i64>();
    let mut d: i64 = 2;
    loop {
        let h = d / 2;
        let dic = classify(&A, d);
        for (&r1, ns1) in dic.iter() {
            let r2 = if r1 <= h { h - r1 } else { h * 3 - r1 };
            if r1 == r2 {
                s += sum_same(ns1) / h
            }
            else if r1 < r2 {
                match dic.get(&r2) {
                    Some(ns2) => { s += sum_diff(ns1, ns2) / h },
                    None      => ()
                }
            }
        }
        
        if !is_remained(&dic, d) {
            break
        }
        
        d <<= 1
    }
    s
}

fn main() {
    let A = read_input();
    println!("{}", F(A))
}