[Python]2進数 エイシング2020D

エイシング2020D

Xは非常に大きな値になるので、単純にシミュレートすることはできません。
ここで、Xのpopcountは高々Nであることから、Xをpopcountで割った値は必ずN未満となります。
また、Xiを反転させたとき、そのような整数のpopcountは(Xのpopcount±1)に等しいです。
したがって、まず1～Nまでのすべての数について、それらをpopcountで割ったあまりに置き換えて0になるまでの操作回数を求めておき、次に、Xi=1であるような各桁について、(Xのpopcount+1)で割ったあまりの総和と(Xのpopcount-1)で割ったあまりの総和を求めておくと、f(Xi)を、Xiが1であれば、(Xのpopcount-1)で割ったあまりの総和-2^iを(Xのpopcount-1)で割ったあまりを求め、これを(Xのpopcount-1)で割ることで、N未満の数を得ることができ、同様に、Xiが0であれば、(Xのpopcount+1)で割ったあまりの総和+2^iを(Xのpopcount+1)で割ったあまりを求め、これを(Xのpopcount+1)で割ることで、N未満の数を得ることができます。
以上より、各Xiについて、0になるまでに必要な操作回数をO(1)で判定できることが分かりました。

ここで、1～Nまでのすべての数について操作回数を求めるための計算量について2*10^5以下の数のpopcountは高々18、18以下の数のpopcountは高々4、4以下の数のpopcountは高々2、2以下の数のpopcountは高々1であり、最大でも5回の操作で0に到達することが分かります。
したがって、O(NlogN)により、1～Nまでのすべての数について操作回数を求めることができることがいえ、この問題をO(NlogN)で解けることが分かりました。
ただし、Xのpopcountが0または1のケースに注意してください。

サンプルコード

n = int(input())
x = input()

# 元のpopcount
original_pop_count = x.count('1')
# 反転して-1count
one_pop_count = original_pop_count - 1
# 反転して+1count
zero_pop_count = original_pop_count + 1

# ±1それぞれにおける余り
if one_pop_count == 0:
  one_mod = 0
else:
  one_mod = int(x, 2) % one_pop_count
zero_mod = int(x, 2) % zero_pop_count

# f の前計算 
f = [0] * 220000
pop_count = [0] * 220000
for i in range(1, 220000):
    pop_count[i] = pop_count[i//2] + i % 2
    f[i] = f[i % pop_count[i]] + 1

# pow2 の前計算
one_pow2 = [1] * 220000
zero_pow2 = [1] * 220000
for i in range(1, 220000):
    if one_pop_count != 0:
        one_pow2[i] = one_pow2[i-1] * 2 % one_pop_count
    zero_pow2[i] = zero_pow2[i-1] * 2 % zero_pop_count

# n-1, n-2, ,, 0
for i in range(n-1, -1, -1):
    if x[n-1-i] == '1':
        if one_pop_count != 0:
            nxt = one_mod
            # - 2^i % 2
            nxt -= one_pow2[i]
            nxt %= one_pop_count
            print(f[nxt] + 1)
        else:
            print(0)
    if x[n-1-i] == '0':
        nxt = zero_mod
        nxt += zero_pow2[i]
        # + 2^i % 2
        nxt %= zero_pop_count
        print(f[nxt] + 1)