Problem p03987 (python)

Fixed submission:
import bisect


class SqrtSet:
    def __init__(self, block_limit=201):
        self.key = []
        self.child = [[]]
        self.block_limit = block_limit

    def search_lower(self, key):
        if key is None:
            return None
        ret = None
        i = bisect.bisect_left(self.key, key)
        if i != 0:
            ret = self.key[i - 1]
        block = self.child[i]
        i = bisect.bisect_left(block, key)
        if i != 0:
            ret = block[i - 1]
        return ret
        
    def search_higher(self, key):
        if key is None:
            return None
        ret = None
        i = bisect.bisect_right(self.key, key)
        if i != len(self.key):
            ret = self.key[i]
        block = self.child[i]
        i = bisect.bisect_right(block, key)
        if i != len(block):
            ret = block[i]
        return ret

    def insert(self, key):
        i = bisect.bisect(self.key, key)
        block = self.child[i]
        bisect.insort(block, key)
        if len(block) == self.block_limit:
            sep = self.block_limit // 2
            self.key.insert(i, block[sep])
            self.child.insert(i + 1, block[sep + 1:])
            self.child[i] = block[:sep]
                
    def dump(self):
        for b in self.child:
            print(len(b), end=" ")
        print("")
            
            
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 = SqrtSet()
    t.insert(-1)
    t.insert(n)
    ans = 0
    for i in range(n):
        right = t.search_higher(idx[i])
        left = t.search_lower(idx[i])
        ans+=(i+1)*(idx[i]-left)*(right-idx[i])
        t.insert(idx[i])
        #print(left,idx[i],right)
    print(ans)


main()

Buggy submission (verdict: Runtime Error):
import sys
input=sys.stdin.readline
def segfunc(x,y):
    return x+y

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]=segfunc(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]=segfunc(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=segfunc(res,seg[p])
        if q&1==1:
            res=segfunc(res,seg[q])
            q-=1
        p=p//2
        q=(q-1)//2
    if p==q:
        res=segfunc(res,seg[p])
    else:
        res=segfunc(segfunc(res,seg[p]),seg[q])
    return res

ide_ele=0
n=int(input())
n+=2
num=2**(n-1).bit_length()
seg=[ide_ele]*2*num

l=list(map(int,input().split()))
p=[]

n-=2
for i in range(n):
    p.append([l[i],i+1])
p.sort(key=lambda x:x[0])
ans=0
update(0,1)
update(n+1,1)
for i in range(n):
    update(p[i][1],1)
    x=query(0,p[i][1])
    ok=n+1
    ng=-1
    while ok-ng>1:
        m=(ok+ng)//2
        if query(0,m)<x-1:
            ng=m
            else:
            ok=m
    left=ok
    ok=n+1
    ng=-1
    while ok-ng>1:
        m=(ok+ng)//2
        if query(0,m)<x-1:
            ng=m
            else:
            ok=m
    right=ok
    ans+=(i+1)*(l[i][1]-left)*(right-l[i][1])
print(ans)