AtCoder Beginner Contest 136 E Max GCD

問題はこれ
E - Max GCD

解説

操作前を  \{A_i\} とし、操作が終わったものを  \{B_i\} とする。
最終的に、全て  r の倍数になったとすると、
 B_i は全て  r で割り切れる。
つまり、\sum_{i=1}^N B_i r で割り切れる。
和は操作で不変なので、
\sum_{i=1}^N A_i  = \sum_{i=1}^N B_i r で割り切れる。
つまり、解の候補は \sum_{i=1}^N A_i の約数に限られる。

ここで、 B_N = \sum_{i=1}^N A_i でそれ以外は  B_i = 0 とすることで、全ての約数で割り切れる  \{B_i\} をつくれることが分かる。
しかし、操作回数の制限  K があるため、解はより複雑なものになる。

ここで、 \sum_{i=1}^N A_i の約数 r を一つ固定する。
r の倍数を考えるの、 \mbox{mod } r で考えれば良い。
 \{t_i\} \{A_i\mbox{ mod }r\} を昇順に並べたものとする。

 \{t_i\} のうち、いくつかを  -t_i して、残りを  +(r-t_i) できるとき、
r の倍数になる数列を  \sum_{\mbox{前半で選んだもの}} t_i 回で実現できる。
これは、小さい方から取ったほうがより小さくできる。

 t_i を小さい方から大きい方に足していくつもりで操作すれば最適になる。

理解の助けになるかもしれないので、いくつか Remark を述べる。

  •  i に対して、 t_i は、 -t_i +(r-t_i) かのいずれか一方の操作が行われる。ただし、d 倍の差はある。
  •  \sum_{\mbox{引き算}}t_i = \sum_{\mbox{足し算}}(r-t_i) となるとき、その操作を実現できる。
  • 上の主張と「小さい方から使えば良い」を使うと、ある地点より左は全て  -t_i の操作が行われて、残りの右の部分は全て  +(r-t_i) の操作が行われたものが、 \{B_i\}  (B_i = 0 \mbox{ mod } r)を実現する。

解答

解答

def factors(N): #約数を全て求める。ただし、順不同
    from collections import deque
    ret = deque()
    middle = int( N**(1/2))
    for i in range(1, middle):
        if N%i == 0:
            ret.append(i)
            ret.append(N//i)
            
    if N%middle == 0:
        ret.append(middle)
        if middle != N//middle:
            ret.append(N//middle)
    return ret

def main():
    N, K = map( int, input().split())
    A = list( map( int, input().split()))
    sa = sum(A)
    R = factors(sa)
    l = len(R)
    E = [0]*l
    ans = 0
    for r in R:
        cnt = 0
        E = list( map( lambda x:x%r, A))
        E.sort()
        s = sum(E)
        now = 0
        if s - E[-1] <= K:
            if ans < r:
                ans = r
            continue
        if s == 0:
            continue
        v = 0
        w = N-1
        cnt = 0
        while w-v != 0:
            if E[v] == 0:
                v += 1
                continue
            if E[w] == 0:
                w -= 1
                continue
            if E[v] + E[w] >= r:
                E[v] -= r - E[w]
                cnt += r - E[w]
                E[w] = 0
                w -= 1
            else:
                E[w] += E[v]
                cnt += E[v]
                E[v] = 0
                v += 1
        if cnt <= K:
            if r > ans:
                ans = r
    print(ans)
if __name__ == '__main__':
    main()


おまけ

計算量は、約数列挙が  \mathcal{O}(\sqrt{N} \max\{A_i\}) で、
各操作は並べ替えが律速段階になって、 \mathcal{O}(N\log(N)) で、
あわせて、
 \mathcal{O}(N\sqrt{N}log(N)\max\{A_i\})
となる。