Problem p02620 (python)

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]
        bitree_add(bit, bit_n, t, d + 2, -1)
        k = bitree_sum(bit, t, d + 1)
        c = bitree_lower_bound(bit, bit_n, bit_d, t, k) - 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
        bitree_add(bit, bit_n, q, d + 2, 1)
        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 = inp[27 + d * 26:27 + d * 27] - 1
    m = inp[27 + d * 27]
    ddd = inp[28 + d * 27::2] - 1
    qqq = inp[29 + d * 27::2] - 1

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

    buf = np.zeros(m, dtype=np.int64)

    for i in range(m):
        diff = update_score(bit, bit_n, bit_d, ccc, sss, ttt, ddd[i], qqq[i])
        ttt[ddd[i]] = qqq[i]
        score += diff
        buf[i] = score

    return buf


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)))

Buggy submission (verdict: Runtime Error):
import os
import sys

import numpy as np


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 + 2
    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 + 1, 1)

    return bit, score


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

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

    ttt[d] = q

    return diff


def solve(inp):
    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 = inp[27 + d * 26:27 + d * 27] - 1
    m = inp[27 + d * 27]
    ddd = inp[28 + d * 27::2] - 1
    qqq = inp[29 + d * 27::2] - 1

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

    buf = np.zeros(m, dtype=np.int64)

    for i in range(m):
        diff = update_score(bit, bit_n, bit_d, ccc, sss, ttt, ddd[i], qqq[i])
        score += diff
        buf[i] = score

    return buf


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

    cc = CC('my_module')
    cc.export('bitree_sum', '(i8[:,:],i8,i8)')(bitree_sum)
    cc.export('bitree_add', '(i8[:,:],i8,i8,i8,i8)')(bitree_add)
    cc.export('bitree_lower_bound', '(i8[:,:],i8,i8,i8,i8)')(bitree_lower_bound)
    cc.export('initial_score', '(i8,i8[:],i8[:,:],i8[:])')(initial_score)
    cc.export('update_score', '(i8[:,:],i8,i8,i8[:],i8[:,:],i8[:],i8,i8)')(update_score)
    cc.export('solve', '(i8[:],)')(solve)
    cc.compile()
    exit()

if os.name == 'posix':
    # noinspection PyUnresolvedReferences
    from my_module import bitree_sum, bitree_add, bitree_lower_bound, initial_score, update_socer, solve
else:
    from numba import njit

    bitree_sum = njit('(i8[:,:],i8,i8)', cache=True)(bitree_sum)
    bitree_add = njit('(i8[:,:],i8,i8,i8,i8)', cache=True)(bitree_add)
    bitree_lower_bound = njit('(i8[:,:],i8,i8,i8,i8)', cache=True)(bitree_lower_bound)
    initial_score = njit('(i8,i8[:],i8[:,:],i8[:])', cache=True)(initial_score)
    update_score = njit('(i8[:,:],i8,i8,i8[:],i8[:,:],i8[:],i8,i8)', cache=True)(update_score)
    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)))