AtCoder ABC #136 E - Max GCD

E - Max GCD

問題

数列に対して、あるものにすると同時に、別の項にをするという操作をK回以下おこなって、取りうる最大のGCDを求めるという問題。

解説

ポイントは操作を行っても数列の総和は変わらず、かつ総和はGCDの倍数になっているはず、というところ。 よって総和の約数を大きい順で条件を満たすか試していけば良い。

どういう条件を調べるかと言うと

1. 各項を約数dで割った余りを求める（あまりは全て足すとdの倍数になるはずである）
2. 余りから0を除外して昇順にならべる
3. 前半を操作、後半を操作するとして最小操作回数を算出
4. 操作回数がK以下なら答え

コードにするとこうなる

fn solve(aa: Vec<isize>, k: usize) -> isize {
    let sum: isize = aa.iter().sum();

    let mut divs = (1..)
        .take_while(|x| x * x <= sum)
        .filter(|x| sum % x == 0)
        .flat_map(|x| {
            if sum / x == x {
                vec![x]
            } else {
                vec![x, sum / x]
            }
        })
        .collect::<Vec<_>>();
    divs.sort();

    for &d in divs.iter().rev() {
        let mut rs: Vec<isize> = aa.iter().map(|&x| x % d).filter(|&x| x != 0).collect();
        rs.sort();

        let rsum = rs.iter().sum::<isize>();
        let need = rs
            .iter()
            .take(rs.len() - (rsum as usize) / (d as usize))
            .sum::<isize>();
        if need <= k as isize {
            return d;
        }
    }

    -1
}

3.を求めるのが難しそうだがコードだと

        let need = rs
            .iter()
            .take(rs.len() - (rsum as usize) / (d as usize))
            .sum::<isize>();

たったコレだけだ。言われてみるとたしかにそのとおりなのだが制限時間内にコレを思いつくのはなかなか難しそうだ。