Problem p02618 (Python)

Buggy submission (status: Time Limit Exceeded):
import os
import sys

import numpy as np


def solve(inp):
    def bitree_sum(bit, t, i):
        s = 0
        while i > 0:
            s += bit[t, i]
            i ^= i & -i
        return s

    def bitree_add(bit, n, t, i, x):
        while i <= n:
            bit[t, i] += x
            i += i & -i

    def bitree_lower_bound(bit, n, d, t, x):
        sum_ = 0
        pos = 0
        for i in range(d, -1, -1):
            k = pos + (1 << i)
            if k <= n and sum_ + bit[t, k] < x:
                sum_ += bit[t, k]
                pos += 1 << i
        return pos + 1

    def initial_score(d, ccc, sss, ttt):
        bit_n = d + 3
        bit = np.zeros((26, bit_n), dtype=np.int64)
        INF = 10**18
        for t in range(26):
            bitree_add(bit, bit_n, t, bit_n - 1, INF)

        last = np.full(26, -1, dtype=np.int64)
        score = 0

        for i in range(d):
            t = ttt[i]
            score += sss[i, t]
            for u in range(26):
                if u == t:
                    continue
                score -= ccc[u] * (i - last[u])
            last[t] = i
            bitree_add(bit, bit_n, t, i + 2, 1)

        return bit, score

    def update_score(bit, bit_n, bit_d, ccc, sss, ttt, d, q):
        diff = 0
        t = ttt[d]
        k = bitree_sum(bit, t, d + 2)
        c = bitree_lower_bound(bit, bit_n, bit_d, t, k - 1) - 2
        e = bitree_lower_bound(bit, bit_n, bit_d, t, k + 1) - 2
        b = ccc[t]
        diff -= b * (d - c) * (e - d)
        diff -= sss[d, t]

        k = bitree_sum(bit, q, d + 2)
        c = bitree_lower_bound(bit, bit_n, bit_d, q, k) - 2
        e = bitree_lower_bound(bit, bit_n, bit_d, q, k + 1) - 2
        b = ccc[q]
        diff += b * (d - c) * (e - d)
        diff += sss[d, q]

        return diff

    d = inp[0]
    ccc = inp[1:27]
    sss = np.zeros((d, 26), dtype=np.int64)
    for r in range(d):
        sss[r] = inp[27 + r * 26 : 27 + (r + 1) * 26]

    ttt = np.random.choice(26, d)

    bit, score = initial_score(d, ccc, sss, ttt)
    bit_n = d + 3
    bit_d = int(np.log2(bit_n))

    for _ in range(10**7):
        cd = np.random.randint(0, d)
        ct = np.random.randint(0, 26)
        diff = update_score(bit, bit_n, bit_d, ccc, sss, ttt, cd, ct)
        if diff > 0:
            score += diff
            bitree_add(bit, bit_n, ttt[cd], cd + 2, -1)
            bitree_add(bit, bit_n, ct, cd + 2, 1)
            ttt[cd] = ct

    return ttt + 1


if sys.argv[-1] == "ONLINE_JUDGE":
    from numba.pycc import CC

    cc = CC("my_module")
    cc.export("solve", "(i8[:],)")(solve)
    cc.compile()
    exit()

if os.name == "posix":
    # noinspection PyUnresolvedReferences
    from my_module import solve
else:
    from numba import njit

    solve = njit("(i8[:],)", cache=True)(solve)
    print("compiled", file=sys.stderr)

inp = np.fromstring(sys.stdin.read(), dtype=np.int64, sep=" ")
ans = solve(inp)
print("\n".join(map(str, ans)))


Fixed submission:
import os
import sys

import numpy as np


def solve(inp):
    def bitree_sum(bit, t, i):
        s = 0
        while i > 0:
            s += bit[t, i]
            i ^= i & -i
        return s

    def bitree_add(bit, n, t, i, x):
        while i <= n:
            bit[t, i] += x
            i += i & -i

    def bitree_lower_bound(bit, n, d, t, x):
        sum_ = 0
        pos = 0
        for i in range(d, -1, -1):
            k = pos + (1 << i)
            if k <= n and sum_ + bit[t, k] < x:
                sum_ += bit[t, k]
                pos += 1 << i
        return pos + 1

    def initial_score(d, ccc, sss, ttt):
        bit_n = d + 3
        bit = np.zeros((26, bit_n), dtype=np.int64)
        INF = 10**18
        for t in range(26):
            bitree_add(bit, bit_n, t, bit_n - 1, INF)

        last = np.full(26, -1, dtype=np.int64)
        score = 0

        for i in range(d):
            t = ttt[i]
            score += sss[i, t]
            for u in range(26):
                if u == t:
                    continue
                score -= ccc[u] * (i - last[u])
            last[t] = i
            bitree_add(bit, bit_n, t, i + 2, 1)

        return bit, score

    def update_score(bit, bit_n, bit_d, ccc, sss, ttt, d, q):
        diff = 0
        t = ttt[d]
        k = bitree_sum(bit, t, d + 2)
        c = bitree_lower_bound(bit, bit_n, bit_d, t, k - 1) - 2
        e = bitree_lower_bound(bit, bit_n, bit_d, t, k + 1) - 2
        b = ccc[t]
        diff -= b * (d - c) * (e - d)
        diff -= sss[d, t]

        k = bitree_sum(bit, q, d + 2)
        c = bitree_lower_bound(bit, bit_n, bit_d, q, k) - 2
        e = bitree_lower_bound(bit, bit_n, bit_d, q, k + 1) - 2
        b = ccc[q]
        diff += b * (d - c) * (e - d)
        diff += sss[d, q]

        return diff

    d = inp[0]
    ccc = inp[1:27]
    sss = np.zeros((d, 26), dtype=np.int64)
    for r in range(d):
        sss[r] = inp[27 + r * 26 : 27 + (r + 1) * 26]

    ttt = np.random.choice(26, d)

    bit, score = initial_score(d, ccc, sss, ttt)
    bit_n = d + 3
    bit_d = int(np.log2(bit_n))

    for _ in range(5 * 10**6):
        cd = np.random.randint(0, d)
        ct = np.random.randint(0, 26)
        diff = update_score(bit, bit_n, bit_d, ccc, sss, ttt, cd, ct)
        if diff > 0:
            score += diff
            bitree_add(bit, bit_n, ttt[cd], cd + 2, -1)
            bitree_add(bit, bit_n, ct, cd + 2, 1)
            ttt[cd] = ct

    return ttt + 1


if sys.argv[-1] == "ONLINE_JUDGE":
    from numba.pycc import CC

    cc = CC("my_module")
    cc.export("solve", "(i8[:],)")(solve)
    cc.compile()
    exit()

if os.name == "posix":
    # noinspection PyUnresolvedReferences
    from my_module import solve
else:
    from numba import njit

    solve = njit("(i8[:],)", cache=True)(solve)
    print("compiled", file=sys.stderr)

inp = np.fromstring(sys.stdin.read(), dtype=np.int64, sep=" ")
ans = solve(inp)
print("\n".join(map(str, ans)))