Problem p02919 (python)

Fixed submission:
class SegmentTree:
    def __init__(self, n, init_value):
        self.n = n
        n2 = 1  # nより大きい2の冪数
        while n2 < n:
            n2 <<= 1
        self.n2 = n2
        self.tree = [[init_value, init_value] for _ in range(n2 << 1)]
        self.ini = init_value

    def update(self, i, x):
        # 注: 同じiへのクエリは複数来ないことを利用して単純化
        i += self.n2
        self.tree[i][0] = x
        while i > 1:
            i >>= 1
            sti = self.tree[i]
            if sti[0] > x:
                sti[1] = sti[0]
                sti[0] = x
            elif sti[1] > x:
                sti[1] = x

    def get(self, a, b):
        result = [self.ini, self.ini]
        q = [(1, 0, self.n2)]

        while q:
            k, l, r = q.pop()

            if a <= l and r <= b:
                stk = self.tree[k]
                if result[0] > stk[0]:
                    result[1] = min(result[0], stk[1])
                    result[0] = stk[0]
                elif result[1] > stk[0]:
                    result[1] = stk[0]
                continue

            m = (l + r) // 2
            k <<= 1
            if a < m and l < b:
                q.append((k, l, m))
            if a < r and l < m:
                q.append((k + 1, m, r))

        return result


def solve(n, ppp):
    str = SegmentTree(n, n)
    right = []
    for i in range(n - 1, -1, -1):
        p = ppp[i]
        right.append(str.get(p, n))
        str.update(p - 1, i)
    right.reverse()

    ans = 0
    stl = SegmentTree(n, 1)
    for i in range(n):
        p = ppp[i]
        l1, l2 = stl.get(p, n)
        l1, l2 = -l1, -l2
        r1, r2 = right[i]
        ans += p * ((l1 - l2) * (r1 - i) + (r2 - r1) * (i - l1))
        stl.update(p - 1, -i)

    return ans


n = int(input())
ppp = list(map(int, input().split()))
print(solve(n, ppp))

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


class RandomizedBinarySearchTree:
    def __init__(self, n):
        n += 10
        self.n = n
        self.children = [[0] * n, [0] * n]
        self.values = [0] * n
        self.counts = [0] * n
        self.indices = list(range(n - 2, -1, -1))
        self.root = -1

    def merge(self, left_root, right_root):
        children = self.children
        counts = self.counts
        li = left_root
        ri = right_root
        stack = []
        while li != -1 and ri != -1:
            if random.randrange(counts[li] + counts[ri]) < counts[li]:
                stack.append((li, 1))
                li = children[1][li]
            else:
                stack.append((ri, 0))
                ri = children[0][ri]

        i = li if li != -1 else ri
        while stack:
            pi, is_right = stack.pop()
            children[is_right][pi] = i
            counts[pi] = counts[children[0][pi]] + counts[children[1][pi]] + 1
            i = pi

        return i

    def split(self, root, x):
        i = root
        lefts, rights = self.children
        values = self.values
        counts = self.counts
        l_stack = []
        r_stack = []
        while i != -1:
            if x < values[i]:
                r_stack.append(i)
                i = lefts[i]
            else:
                l_stack.append(i)
                i = rights[i]

        li, ri = -1, -1
        while l_stack:
            pi = l_stack.pop()
            rights[pi] = li
            counts[pi] = counts[lefts[pi]] + counts[li] + 1
            li = pi
        while r_stack:
            pi = r_stack.pop()
            lefts[pi] = ri
            counts[pi] = counts[ri] + counts[rights[pi]] + 1
            ri = pi

        return li, ri

    def insert(self, x):
        ni = self.indices.pop()
        self.children[0][ni] = self.children[1][ni] = -1
        self.values[ni] = x
        self.counts[ni] = 1

        li, ri = self.split(self.root, x)
        self.root = self.merge(self.merge(li, ni), ri)

    def delete(self, x):
        li, mri = self.split(self.root, x - 1)
        mi, ri = self.split(mri, x)

        if mi == -1:
            self.root = self.merge(li, ri)
            return

        self.indices.append(mi)
        self.root = self.merge(li, ri)
        return

    def upper_bound(self, x, default=-1):
        i = self.root
        lefts, rights = self.children
        values = self.values
        counts = self.counts
        y = default
        c = counts[i]
        j = 0
        while i != -1:
            if x < values[i]:
                y = values[i]
                c = j + counts[lefts[i]]
                i = lefts[i]
            else:
                j += counts[lefts[i]] + 1
                i = rights[i]
        return y, c

    def lower_bound(self, x, default=-1):
        i = self.root
        lefts, rights = self.children
        values = self.values
        counts = self.counts
        y = default
        c = counts[i]
        j = 0
        while i != -1:
            if x <= values[i]:
                y = values[i]
                c = j + counts[lefts[i]]
                i = lefts[i]
            else:
                j += counts[lefts[i]] + 1
                i = rights[i]
        return y, c

    def get_k_th(self, k, default=-1):
        i = self.root
        children = self.children
        values = self.values
        counts = self.counts
        if counts[i] <= k:
            return default
        j = k
        while i != -1:
            left_count = counts[children[0][i]]
            if left_count == j:
                return values[i]
            elif left_count > j:
                i = children[0][i]
            else:
                j -= left_count + 1
                i = children[1][i]
        return default

    def debug_print(self):
        print('Lefts ', self.children[0])
        print('Rights', self.children[1])
        print('Values', self.values)
        print('Counts', self.counts)
        self._debug_print(self.root, 0)

    def _debug_print(self, i, depth):
        if i != -1:
            self._debug_print(self.children[0][i], depth + 1)
            print('      ' * depth, self.values[i], self.counts[i])
            self._debug_print(self.children[1][i], depth + 1)


q = int(input())
rbst = RandomizedBinarySearchTree(q)
for line in sys.stdin:
    t, x = map(int, line.split())
    if t == 1:
        rbst.insert(x)
    else:
        y = rbst.get_k_th(x - 1)
        print(y)
        rbst.delete(y)

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

import numpy as np


# Global variables
# N is large enough size to store nodes at one time.
# children: ndarray(N, 2)  current [left / right] child node index
# values:   ndarray(N)     -1 is represents EOT
# colors:   ndarray(N)     0:black 1:red
# counts:   ndarray(N)     number of nodes in subtree
# available: set           available indices, set(range(N)) at initial
# root:      int           root index, 0 at initial


def solve(inp):
    RBT_SIZE = 20
    RBT_EOT = -2
    RBT_AVAILABLE_INDICES = set(range(RBT_SIZE - 1))
    RBT_CHILDREN = np.zeros((RBT_SIZE, 2), dtype=np.int64)
    RBT_VALUES = np.zeros(RBT_SIZE, dtype=np.int64)
    RBT_COLORS = np.zeros(RBT_SIZE, dtype=np.int64)
    RBT_COUNTS = np.zeros(RBT_SIZE, dtype=np.int64)
    RBT_ROOT = -1
    RBT_CHILDREN[RBT_ROOT, 0] = RBT_CHILDREN[RBT_ROOT, 1] = RBT_ROOT

    def _rbt_rotate(i, is_right):
        """
        Rotate node i to right if is_right=1 and left if 0.
        Return the index of new node that replaces node i.
        """
        r = is_right
        ci = RBT_CHILDREN[i, r ^ 1]
        RBT_CHILDREN[i, r ^ 1] = RBT_CHILDREN[ci, r]
        RBT_CHILDREN[ci, r] = i
        RBT_COLORS[ci] = RBT_COLORS[i]
        RBT_COLORS[i] = 1
        child_count = RBT_COUNTS[ci]
        RBT_COUNTS[ci] = RBT_COUNTS[i]
        RBT_COUNTS[i] -= child_count - RBT_COUNTS[RBT_CHILDREN[i, r ^ 1]]
        return ci

    def _rbt_balance_insert(i, is_right):
        flag = True
        r = is_right
        if RBT_COLORS[RBT_CHILDREN[RBT_CHILDREN[i, r], r ^ 1]] == 1:
            RBT_CHILDREN[i, r] = _rbt_rotate(RBT_CHILDREN[i, r], r)
        if RBT_COLORS[RBT_CHILDREN[RBT_CHILDREN[i, r], r]] == 1:
            if RBT_COLORS[RBT_CHILDREN[i, r ^ 1]] == 1:
                RBT_COLORS[i] = 1
                RBT_COLORS[RBT_CHILDREN[i, r]] = RBT_COLORS[RBT_CHILDREN[i, r ^ 1]] = 0
                flag = False
            else:
                i = _rbt_rotate(i, r ^ 1)
        return i, flag

    def _rbt_balance_delete(i, is_right):
        flag = True
        stack = [(i, is_right)]
        while stack:
            i, r = stack.pop()
            if RBT_COLORS[RBT_CHILDREN[RBT_CHILDREN[i, r ^ 1], r]] == 0 and \
                    RBT_COLORS[RBT_CHILDREN[RBT_CHILDREN[i, r ^ 1], r ^ 1]] == 0:
                if RBT_COLORS[RBT_CHILDREN[i, r ^ 1]] == 0:
                    RBT_COLORS[RBT_CHILDREN[i, r ^ 1]] = 1
                    if RBT_COLORS[i] == 0:
                        flag = False
                    RBT_COLORS[i] = 0
                    break
                else:
                    i = _rbt_rotate(i, r)
                    stack.append((i, r))
                    stack.append((RBT_CHILDREN[i, r], r))
                    continue
            else:
                if RBT_COLORS[RBT_CHILDREN[RBT_CHILDREN[i, r ^ 1], r]] == 1:
                    RBT_CHILDREN[i, r ^ 1] = _rbt_rotate(RBT_CHILDREN[i, r ^ 1], r ^ 1)
                i = _rbt_rotate(i, r)
                RBT_COLORS[RBT_CHILDREN[i, r]] = 0
                RBT_COLORS[RBT_CHILDREN[i, r ^ 1]] = 0
                break

        while stack:
            pi, r = stack.pop()
            RBT_CHILDREN[pi, r] = i
            i = pi

        return i, flag

    def _rbt_get_min_in_subtree(i, stack):
        while RBT_CHILDREN[i, 0] != -1:
            stack.append((i, 0))
            i = RBT_CHILDREN[i, 0]
        return i

    def rbt_insert(root, x):
        """
        Insert value x.
        Return new root index
        """

        stack = []
        i = root
        while i != -1:
            to_right = 1 if x >= RBT_VALUES[i] else 0
            stack.append((i, to_right))
            i = RBT_CHILDREN[i, to_right]

        for pi, _ in stack:
            RBT_COUNTS[pi] += 1

        i = RBT_AVAILABLE_INDICES.pop()
        RBT_CHILDREN[i, 0] = -1
        RBT_CHILDREN[i, 1] = -1
        RBT_VALUES[i] = x
        RBT_COLORS[i] = 1
        RBT_COUNTS[i] = 1

        if stack:
            pi, r = stack[-1]
            RBT_CHILDREN[pi, r] = i

        while stack:
            pi, is_right = stack.pop()
            if RBT_COLORS[pi] == 1:
                i = pi
                continue
            RBT_CHILDREN[pi, is_right] = i
            i, flag = _rbt_balance_insert(pi, is_right)

            if stack and flag:
                pi, is_right = stack.pop()
                RBT_CHILDREN[pi, is_right] = i
                break

        else:
            RBT_COLORS[i] = 0
            return i

        return root

    def rbt_delete(root, x):
        """
        Delete value x.
        Return new root index
        """
        i = root
        stack = []
        while i != -1 and RBT_VALUES[i] != x:
            is_right = 1 if RBT_VALUES[i] < x else 0
            stack.append((i, is_right))
            i = RBT_CHILDREN[i, is_right]

        if i == -1:
            return root

        if RBT_CHILDREN[i, 0] != -1 and RBT_CHILDREN[i, 1] != -1:
            stack.append((i, 1))
            mi = _rbt_get_min_in_subtree(RBT_CHILDREN[i, 1], stack)
            RBT_COLORS[i] = RBT_COLORS[mi]
            i = mi

        for pi, _ in stack:
            RBT_COUNTS[pi] -= 1

        pi, is_right = stack[-1]
        if RBT_CHILDREN[i, 0] == -1:
            RBT_CHILDREN[pi, is_right] = RBT_CHILDREN[i, 1]
            RBT_COLORS[RBT_CHILDREN[i, 1]] = 0
            if RBT_CHILDREN[i, 1] != -1 or RBT_COLORS[i] == 1:
                return root
        elif RBT_CHILDREN[i, 1] == -1:
            RBT_CHILDREN[pi, is_right] = RBT_CHILDREN[i, 0]
            RBT_COLORS[RBT_CHILDREN[i, 0]] = 0
            return root

        while stack:
            pi, is_right = stack.pop()
            i, flag = _rbt_balance_delete(pi, is_right)

            if stack and flag:
                pi, is_right = stack.pop()
                RBT_CHILDREN[pi, is_right] = i
                break

        else:
            RBT_COLORS[i] = 0
            return i

        return root

    def rbt_upper_bound(root, x):
        i = root
        y = RBT_EOT
        c = RBT_COUNTS[i]
        j = 0
        while i != -1:
            if x < RBT_VALUES[i]:
                y = RBT_VALUES[i]
                c = j + RBT_COUNTS[RBT_CHILDREN[i, 0]]
                i = RBT_CHILDREN[i, 0]
            else:
                j += RBT_COUNTS[RBT_CHILDREN[i, 0]] + 1
                i = RBT_CHILDREN[i, 1]
        return y, c

    def rbt_lower_bound(root, x):
        i = root
        y = RBT_EOT
        c = RBT_COUNTS[i]
        j = 0
        while i != -1:
            if x <= RBT_VALUES[i]:
                y = RBT_VALUES[i]
                c = j + RBT_COUNTS[RBT_CHILDREN[i, 0]]
                i = RBT_CHILDREN[i, 0]
            else:
                j += RBT_COUNTS[RBT_CHILDREN[i, 0]] + 1
                i = RBT_CHILDREN[i, 1]
        return y, c

    def rbt_get_k_th(root, k):
        i = root
        if RBT_COUNTS[i] <= k:
            return RBT_EOT
        j = k
        while i != -1:
            left_count = RBT_COUNTS[RBT_CHILDREN[i, 0]]
            if left_count == j:
                return RBT_VALUES[i]
            elif left_count > j:
                i = RBT_CHILDREN[i, 0]
            else:
                j -= left_count + 1
                i = RBT_CHILDREN[i, 1]
        return RBT_EOT

    n = inp[0]
    ppp = inp[1:]
    idx = np.argsort(ppp)[::-1]

    RBT_ROOT = rbt_insert(RBT_ROOT, -1)
    RBT_ROOT = rbt_insert(RBT_ROOT, -1)
    RBT_ROOT = rbt_insert(RBT_ROOT, n)
    RBT_ROOT = rbt_insert(RBT_ROOT, n)

    ans = 0
    for i in idx:
        p = ppp[i]
        r1, j = rbt_upper_bound(RBT_ROOT, i)
        r2 = rbt_get_k_th(RBT_ROOT, j + 1)
        l1 = rbt_get_k_th(RBT_ROOT, j - 1)
        l2 = rbt_get_k_th(RBT_ROOT, j - 2)
        ans += p * ((l1 - l2) * (r1 - i) + (r2 - r1) * (i - l1))
        RBT_ROOT = rbt_insert(RBT_ROOT, i)

    return ans


def debug_print(children, values, colors, counts, root):
    print(children.T)
    print(values)
    print(colors)
    print(counts)
    _debug_print(children, values, colors, counts, root, 0)


def _debug_print(children, values, colors, counts, i, depth):
    if i != -1:
        _debug_print(children, values, colors, counts, children[i, 0], depth + 1)
        print('      ' * depth, 'BR'[colors[i]], values[i], counts[i])
        _debug_print(children, values, colors, counts, children[i, 1], depth + 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)
    pass

inp = np.fromstring(sys.stdin.read(), dtype=np.int64, sep=' ')
ans = solve(inp)
print(ans)

Buggy submission (verdict: Runtime Error):
class SegmentTree:
    """
    値をlistで持つSegmentTree
    update(i, x): i(0-indexed)をxで更新
    get(a,b): [a, b)を取得
    """

    def __init__(self, n, ini, merge):
        self.n = n
        n2 = 1  # nより大きい2の冪数
        while n2 < n:
            n2 <<= 1
        self.n2 = n2
        self.tree = [ini.copy() for _ in range(n2 << 1)]
        self.ini = ini
        self.merge = merge

    def update(self, i, x):
        i += self.n2
        self.tree[i] = self.merge(x, self.tree[i])
        while i > 1:
            self.tree[i >> 1] = x = self.merge(x, self.tree[i ^ 1])
            i >>= 1

    def get(self, a, b):
        return self._get(a, b, 1, 0, self.n2)

    def _get(self, a, b, k, l, r):
        """
        [a, b)の最小値を得る内部関数
        :param k: 現在調べている区間のtree内index
        :param l, r: kが表す区間の左右端index [l, r)
        :return: kが表す区間と[a, b)の共通区間内での値
        """
        # 範囲外ならINF
        if r <= a or b <= l:
            return self.ini
            # [a,b)が完全に[l,r)を包含するならtree[k]の値を採用
        if a <= l and r <= b:
            return self.tree[k]
        # 一部だけ範囲内なら2つに分けて再帰的に調査
        m = (l + r) // 2
        return self.merge(self._get(a, b, k << 1, l, m), self._get(a, b, (k << 1) + 1, m, r))


def solve(n, ppp):
    def merge_min(x, y):
        if x[0] <= y[0]:
            if x[1] <= y[0]:
                return x.copy()
            return [x[0], y[0]]
        if x[0] <= y[1]:
            return [y[0], x[0]]
        return y.copy()

    def merge_max(x, y):
        if x[0] >= y[0]:
            if x[1] >= y[0]:
                return x.copy()
            return [x[0], y[0]]
        if x[0] >= y[1]:
            return [y[0], x[0]]
        return y.copy()

    stl = SegmentTree(n, [-1, -1], merge_max)
    left1 = [-1] * n
    left2 = [-1] * n
    for i, p in enumerate(ppp):
        res = stl.get(p, n)
        left1[i] = res[0]
        left2[i] = res[1]
        stl.update(p - 1, [i, -1])

    str = SegmentTree(n, [n, n], merge_min)
    right1 = [n] * n
    right2 = [n] * n
    for i, p in enumerate(ppp[::-1]):
        i = n - i - 1
        res = str.get(p, n)
        right1[i] = res[0]
        right2[i] = res[1]
        str.update(p - 1, [i, n])
    # print(left1)
    # print(left2)
    # print(right1)
    # print(right2)

    ans = 0
    for i, l1, l2, r1, r2 in zip(range(n), left1, left2, right1, right2):
        ans += ppp[i] * ((l1 - l2) * (r1 - i) + (r2 - r1) * (i - l1))
    return ans


n = int(input())
ppp = list(map(int, input().split()))
print(solve(n, ppp))