Problem p02919 (python)

Fixed submission:
n = int(input())
p = list(map(int,input().split()))

def init(init_val):
    #set_val
    for i in range(n):
        seg[i+num-1]=init_val[i]    
    #built
    for i in range(num-2,-1,-1) :
        seg[i]=max(seg[2*i+1],seg[2*i+2]) 
    
def update(k,x):
    k += num-1
    seg[k] = x
    while k+1:
        k = (k-1)//2
        seg[k] = max(seg[k*2+1],seg[k*2+2])
    
def query(p,q):
    if q<=p:
        return ide_ele
    p += num-1
    q += num-2
    res=ide_ele
    while q-p>1:
        if p&1 == 0:
            res = max(res,seg[p])
        if q&1 == 1:
            res = max(res,seg[q])
            q -= 1
        p = p//2
        q = (q-1)//2
    if p == q:
        res = max(res,seg[p])
    else:
        res = max(max(res,seg[p]),seg[q])
    return res

#####単位元######
ide_ele = -1

#num:n以上の最小の2のべき乗
num =2**(n).bit_length()
seg=[ide_ele]*2*num

res = 0
s = [0]*(n+1)

update(p[0],0)
for i in range(1,n):
    k1 = query(p[i],n+1)
    if k1 != -1:
        update(p[k1],-1)
        k2 = query(p[i],n+1)
        s[i] = p[i]*(k1-k2) + s[k1]
        res += s[i]
        update(p[k1],k1)
    update(p[i],i)


ide_ele = -n
seg=[ide_ele]*2*num

s = [0]*(n+1)

update(p[-1],1-n)
for i in range(n-2,-1,-1):
    k1 = query(p[i],n+1)
    if k1 != -n:
        update(p[-k1],-n)
        k2 = query(p[i],n+1)
        s[i] = p[i]*(k1-k2)+s[-k1]
        res += s[i]
        update(p[-k1],k1)
    update(p[i],-i)

print(res)

Buggy submission (verdict: Runtime Error):
from time import time
from statistics import median
from bisect import bisect_left, bisect_right
from random import random, sample


class TwoLayerSet:
    def __init__(self, p):
        self.top = []
        self.bottom = [[]]
        self.p = p

    def insert(self, key):
        (top, bottom) = (self.top, self.bottom)
        top_i = bisect_left(top, key)
        if top_i != len(top) and key == top[top_i]:
            return False
        block = bottom[top_i]
        block_i = bisect_left(block, key)
        if block_i != len(block) and key == block[block_i]:
            return False
        if random() < self.p:
            top.insert(top_i, key)
            bottom[top_i] = block[block_i:]
            bottom.insert(top_i, block[:block_i])
        else:
            block.insert(block_i, key)
        return True

    def erase(self, key):
        (top, bottom) = (self.top, self.bottom)
        top_i = bisect_left(top, key)
        if top_i != len(top) and key == top[top_i]:
            top.pop(top_i)
            bottom[top_i].extend(bottom.pop(top_i+1))
            return True
        block = bottom[top_i]
        block_i = bisect_left(block, key)
        if block_i != len(block) and key == block[block_i]:
            block.pop(block_i)
            return True
        return False

    def less(self, key):
        if key is None:
            return None
        top = self.top
        top_i = bisect_left(top, key)
        block = self.bottom[top_i]
        block_i = bisect_left(block, key)
        if block_i != 0:
            return block[block_i - 1]
        if top_i != 0:
            return top[top_i - 1]
        return None

    def less_equal(self, key):
        top = self.top
        top_i = bisect_right(top, key)
        block = self.bottom[top_i]
        block_i = bisect_right(block, key)
        if block_i != 0:
            return block[block_i - 1]
        if top_i != 0:
            return top[top_i - 1]
        return None

    def greater(self, key):
        if key is None:
            return None
        top = self.top
        top_i = bisect_right(top, key)
        block = self.bottom[top_i]
        block_i = bisect_right(block, key)
        if block_i != len(block):
            return block[block_i]
        if top_i != len(top):
            return top[top_i]
        return None

    def greater_equal(self, key):
        top = self.top
        top_i = bisect_left(top, key)
        block = self.bottom[top_i]
        block_i = bisect_left(block, key)
        if block_i != len(block):
            return block[block_i]
        if top_i != len(top):
            return top[top_i]
        return None


def main():
    n = int(input())
    p = list(map(int, input().split()))
    idx = [0] * n
    for i in range(0, n):
        idx[i] = i
    idx.sort(key=lambda i: - p[i])
    t = TwoLayerSet(1/100)
    t.insert(-1)
    t.insert(n)
    ans = 0
    for i in idx:
        nex = t.greater(i)
        nexnex = t.greater(nex)
        pre = t.less(i)
        prepre = t.less(pre)
        if prepre != None:
            ans += p[i] * (pre - prepre) * (nex - i)
        if nexnex != None:
            ans += p[i] * (i - pre) * (nexnex - nex)
        t.insert(i)
    print(ans)


main()