Problem p02919 (Python)

Buggy submission (status: 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)


Fixed submission:
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 = 100010
    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)