Problem p02609 (Python)

Buggy submission (status: Runtime Error):
def main():
    N = int(input())
    (*X,) = map(int, input())

    memo = [-1] * 21

    def f(x):
        ret = 0
        while x > 0:
            x = x % bin(x)[2:].count("1")
            ret += 1
        return ret

    for x in range(21):
        memo[x] = f(x)

    cnt = X.count(1)

    pow2_0_1 = [1]
    pow2_1_0 = [1]

    t_0_1 = 1
    t_1_0 = 1
    for i in range(N):
        t_0_1 = t_0_1 * 2 % (cnt + 1)
        pow2_0_1.append(t_0_1)

        t_1_0 = t_1_0 * 2 % (cnt - 1)
        pow2_1_0.append(t_1_0)

    chg_0_1 = 0
    chg_1_0 = 0
    i = 0
    for d in reversed(X):
        if d:
            chg_0_1 = (chg_0_1 + pow2_0_1[i]) % (cnt + 1)
            chg_1_0 = (chg_1_0 + pow2_1_0[i]) % (cnt - 1)
        i += 1

    ans = []
    i = 0
    for d in reversed(X):
        if d:
            t = (chg_1_0 - pow2_1_0[i]) % (cnt - 1)
        else:
            t = (chg_0_1 + pow2_0_1[i]) % (cnt + 1)
        i += 1

        add_v = 1
        while t > 20 or memo[t] == -1:
            t = t % bin(t)[2:].count("1")
            add_v += 1
        ans.append(memo[t] + add_v)

    ans.reverse()

    print(*ans, sep="\n")


if __name__ == "__main__":
    main()

# import sys
# input = sys.stdin.readline
#
# sys.setrecursionlimit(10 ** 7)
#
# (int(x)-1 for x in input().split())
# rstrip()
#
# def binary_search(*, ok, ng, func):
#     while abs(ok - ng) > 1:
#         mid = (ok + ng) // 2
#         if func(mid):
#             ok = mid
#         else:
#             ng = mid
#     return ok


Fixed submission:
def main():
    N = int(input())
    (*X,) = map(int, input())

    memo = [-1] * 21

    def f(x):
        ret = 0
        while x > 0:
            x = x % bin(x)[2:].count("1")
            ret += 1
        return ret

    for x in range(21):
        memo[x] = f(x)

    cnt = X.count(1)

    if cnt == 1:
        ans = []
        i = 0
        for d in reversed(X):
            if d:
                ans.append(0)
                continue
            if i == 0 or X[-1]:
                ans.append(2)
            else:
                ans.append(1)
            i += 1

        ans.reverse()

        print(*ans, sep="\n")
        return

    pow2_0_1 = [1]
    pow2_1_0 = [1]

    t_0_1 = 1
    t_1_0 = 1
    for i in range(N):
        t_0_1 = t_0_1 * 2 % (cnt + 1)
        pow2_0_1.append(t_0_1)

        t_1_0 = t_1_0 * 2 % (cnt - 1)
        pow2_1_0.append(t_1_0)

    chg_0_1 = 0
    chg_1_0 = 0
    i = 0
    for d in reversed(X):
        if d:
            chg_0_1 = (chg_0_1 + pow2_0_1[i]) % (cnt + 1)
            chg_1_0 = (chg_1_0 + pow2_1_0[i]) % (cnt - 1)
        i += 1

    ans = []
    i = 0
    for d in reversed(X):
        if d:
            t = (chg_1_0 - pow2_1_0[i]) % (cnt - 1)
        else:
            t = (chg_0_1 + pow2_0_1[i]) % (cnt + 1)
        i += 1

        add_v = 1
        while t > 20 or memo[t] == -1:
            t = t % bin(t)[2:].count("1")
            add_v += 1
        ans.append(memo[t] + add_v)

    ans.reverse()

    print(*ans, sep="\n")


if __name__ == "__main__":
    main()

# import sys
# input = sys.stdin.readline
#
# sys.setrecursionlimit(10 ** 7)
#
# (int(x)-1 for x in input().split())
# rstrip()
#
# def binary_search(*, ok, ng, func):
#     while abs(ok - ng) > 1:
#         mid = (ok + ng) // 2
#         if func(mid):
#             ok = mid
#         else:
#             ng = mid
#     return ok