AtCoder ABC 136 E - Max GCD (Go) - あるエンジニアのAtCoder奮闘記

irisruneです。問題の解説など試行錯誤していますが、アドバイスや感想などコメントを貰えると嬉しいです。

問題

あることに気付かないとどうにも方針が立てられないので、かなりパズル要素が強い問題だと思います。

package main

import (
    "bufio"
    "fmt"
    "os"
    "sort"
    "strconv"
)

var sc *bufio.Scanner

func nextInt() int {
    sc.Scan()
    i, e := strconv.Atoi(sc.Text())
    if e != nil {
        panic(e)
    }
    return i
}

func main() {
    sc = bufio.NewScanner(os.Stdin)
    sc.Split(bufio.ScanWords)
    n, k := nextInt(), nextInt()

    aArray := make([]int, n)
    aSum := 0
    for i := range aArray {
        aArray[i] = nextInt()
        aSum += aArray[i]
    }

    var div []int
    for i := 1; i*i <= aSum; i++ {
        if aSum%i == 0 {
            div = append(div, i)
            if i != aSum/i {
                div = append(div, aSum/i)
            }
        }
    }
    sort.Sort(sort.Reverse(sort.IntSlice(div)))

    for _, d := range div {
        mod := make([]int, n)
        for i, a := range aArray {
            mod[i] = a % d
        }
        sort.Sort(sort.IntSlice(mod))
        modSum := 0
        for _, m := range mod {
            modSum += m
        }
        if modSum == 0 {
            fmt.Println(d)
            return
        }
        modDiff := 0
        for i := n - 1; modSum > modDiff; i-- {
            modSum -= mod[i]
            modDiff += d - mod[i]
        }
        if modSum == modDiff && modSum <= k {
            fmt.Println(d)
            return
        }
    }
}

何回操作を行っても $A_i$ の和は不変であること、そして答えとなる値は $A_i$ の和の約数しかあり得ないということに気付くことがスタートラインです。（ $K$ を無限大とおいた上では）十分条件であることは簡単に証明できますが、必要条件であることを証明するのは少々骨が折れそうです。~~公式解説でも証明されていないので~~今回は証明なしで進めていきます。

$A_i$ の和の約数に答えの候補を絞りこんだところで、実際に答えとなる値を求める方法について考えます。求めるゴールとしては $A_i$ が全て答えの倍数になるよう操作を行うことを考えることになります。

ここで $K$ を無限大とおいた場合、すべての答えの候補に対してこのような操作は必ず行うことができます。というのも、 $A_i$ のうち1つを除いて全て $0$ になるよう操作を行えば、除いた1つの値は $A_i$ の和となり、すべての $A_i$ の和の約数は条件を満たすことになるからです。そういうわけで、実際には操作回数が $K$ 以下であるかどうかをチェックすればよいことになります（上記コードでは操作が行えるか自体もチェックしていますが）。

ではどのようにチェックすればよいかというと、答えの候補の値で $A_i$ それぞれを割った剰余に注目し、剰余が小さい方から一定の個数の和を操作回数とします。操作回数に含めなかった分については、答えの候補の値から剰余を引いた値の和を負の操作回数とし、2つの操作回数の絶対値が一致する値を真の操作回数として求めます。あとは真の操作回数が $K$ 以下であれば答えとして成立するということになります（出力するのは最大値のみということに注意）。

公式解説では累積和などを用いればよいとありますが、おそらく両方向の累積和を計算する必要がありやや複雑化するため今回は尺取り法で求めました。計算量については変わらず、 $A_i\leq10^6$ なので $O(N\sqrt{10^6N}\log N)$ となります（定数が入るのはオーダー表記としておかしいですが、無視できない値なので大目に見てください）。