AtCoder ABC 134 D - Preparing Boxes (Python) - あるエンジニアのAtCoder奮闘記

irisruneです。とりあえずABC134のD,Eはどうにかなりました。

問題

計算量の見積りが大変難しい問題です。また、逆から解くというアプローチが求められます。

import sys


def input():
    return sys.stdin.readline().rstrip()


def main():
    n = int(input())
    bits = [int(e) for e in input().split()]
    ans = set()
    for i in reversed(range(n)):
        for j in range(i + (i + 1), n, i + 1):
            bits[i] ^= bits[j]
        if bits[i] == 1:
            ans.add(i + 1)
    print(len(ans))
    print(*ans)


main()

まず1番目の箱から順にボールが入っているか判定することを考えてみます。例えば、 $a_i(i\geq1)$ の排他的論理和を計算するなどして、1番目の箱にボールが入っているかどうかを判定できるか考えてみます。すると、

3
1 0 1

という入力では1番目の箱にボールは入らないですが（3番目のみ入る）、

4
1 0 0 1

という入力では1番目の箱にボールが入ることになります（1,2,4番目に入る）。他には、

5
1 0 0 0 1

という入力では1番目の箱にボールは入らず（5番目のみ入る）、

6
1 0 0 0 0 1

という入力でも1番目の箱にボールは入りません（2,3,6番目に入る）。以上から、これらの場合分けを $a_i$ からそのまま行うのはおそらく不可能と推測できます。

つまり、1番目の箱にボールが入っているかどうかは、1の倍数の書かれた箱（つまり2番目以降すべての箱）にボールが入っているかどうかがわからないと判定できないといえるでしょう。

2番目以降の箱についても同様に考えると、結局 $N$ 番目の箱から逆順にボールが入っているかどうかを判定するアプローチを取るのがよさそうです。

あとは、 $i$ 番目の箱にボールが入っているかどうかの判定を $i$ の倍数が書かれた箱にボールが入っているかチェックしながら行うことで愚直に解くことができます。その計算量は $O(N/1+N/2+\ldots+N/N)$ となり高速には見えませんが、

$$ \begin{align} & N/1+N/2+N/3+N/4+N/5+N/6+N/7+\ldots+N/N \\ \leq & N/1+N/2+N/2+N/4+N/4+N/4+N/4+\ldots+(N/2^{\lfloor\log_2N\rfloor}) \\ \leq & N+N+N+\ldots+N \\ = & N\lfloor\log_2N\rfloor \end{align} $$

と計算することができるため結局計算量は $O(N\log N)$ で表せ、十分高速であることがわかります。

ちなみに、上記コードではボールが入るかどうかのフラグ管理を $a_i$ を書き換えることで行っているためアルゴリズムが読み取りにくくなっていると思います。途中計算とは別に出力用にボールが入る箱のセットを用意していたのはいいのですが、途中計算用のリストも別に作った方がよかったですね。

雑記

計算量に真面目に向き合った結果、前回よりも数学要素が強くなった気がします。