Problem p02919 (python)

Fixed submission:
import sys

stdin = sys.stdin

ni = lambda: int(ns())
nl = lambda: list(map(int, stdin.readline().split()))
nm = lambda: map(int, stdin.readline().split())
ns = lambda: stdin.readline().rstrip()

n = ni()
p = nl()
q = [0]*(n+1)
for i in range(n):
    q[p[i]] = i+1
    
ans = 0
l = [0] + [i for i in range(n+1)]
r = [i+1 for i in range(n+1)] + [n+1]

for i in range(1,n+1):
    v = q[i]
    l1,r1 = l[v],r[v]
    l2,r2 = l[l1],r[r1]
    ans += i*((v-l1)*(r2-r1) + (r1-v)*(l1-l2))
    l[r1] = l1
    r[l1] = r1
  
print(ans)

Buggy submission (verdict: Runtime Error):
class B():
    def __init__(self,N):
        self.N = N
        self.node = [0]*(self.N+1)
        self.cnt = 0

    def add(self,x): # 要素 x を追加
        self.cnt += 1
        while x <= self.N:
            self.node[x] += 1
            x += x & -x

    def delete(self,x): # 要素 x を削除
        self.cnt -= 1
        while x <= self.N:
            self.node[x] -= 1
            x += x & -x

    def count(self,x): # x以下の要素数
        tmp = 0
        while x > 0:
            tmp += self.node[x]
            x -= x & -x
        return tmp

    def get_max(self):
        return self.get_ith(self.cnt)

    def get_ith(self,i): # i 番目に小さい要素を取得
        NG = -1
        OK = self.N+1
        while OK-NG > 1:
            mid = (OK+NG)//2
            #print(OK,NG,self.count(mid))
            if self.count(mid) >= i:
                OK = mid
            else:
                NG = mid
        return OK

    def search_low(self,x): # x より小さい最大の要素
        xcn = self.count(x)
        NG = -1
        OK = x
        while OK-NG > 1:
            mid = (OK+NG)//2
            #print(OK,NG,self.count(mid))
            if self.count(mid) >= xcn:
                OK = mid
            else:
                NG = mid
        return OK

    def search_high(self,x): # x より大きい最小の要素
        xcn = self.count(x)
        NG = x
        OK = self.N + 1
        while OK-NG > 1:
            mid = (OK+NG)//2
            #print(OK,NG,self.count(mid))
            if self.count(mid) > xcn:
                OK = mid
            else:
                NG = mid
        return OK

import sys

stdin = sys.stdin

ni = lambda: int(ns())
nl = lambda: list(map(int, stdin.readline().split()))
nm = lambda: map(int, stdin.readline().split())
ns = lambda: stdin.readline().rstrip()

n = ni()
p = nl()
q = [0]*(n+1)
for i in range(n):
    q[p[i]] = i+1
ans = 0
b = B(n)
b.add(q[-1])
for i in range(n-1,0,-1):
  v = q[i]
  l1 = b.search_low(v)
  r1 = b.search_high(v)
  l2 = b.search_low(l1)
  r2 = b.search_high(r1)
  ans += i*((r1-v)*(l1-l2) + (v-l1)*(r2-r1))
  b.add(v)
print(ans)

Buggy submission (verdict: Runtime Error):
class Avltree:
    """ Avl木を表現するクラス """
    def __init__(self):
        self.avl = [None,None,None,0]

    def node(self,key):
        return [None,key,None,0]

    # search(x) xが存在するならTrueをしないならFalseを返す
    # 使用法 Instance.search(x)
    def search(self,key,avl=None):
        if avl == None:
            avl = self.avl
        if key == avl[1]:
            return True
        elif key < avl[1]:
            if avl[0] == None:
                return False
            else:
                return self.search(key,avl[0])
        elif avl[1] < key:
            if avl[2] == None:
                return False
            else:
                return self.search(key,avl[2])
    
    # key未満で最大の要素を検索する。なければNoneを返す。
    # 使用法 Instance.search_lower(x)
    def search_lower(self,key,key_lower=None,avl=0):
        if avl == 0:
            avl = self.avl
        if avl == None:
            return key_lower
        elif key < avl[1]:
            if avl[0] == None:
                return key_lower
            else:
                return self.search_lower(key,key_lower,avl[0])
        elif avl[1] < key:
            key_lower = avl[1]
            if avl[2] == None:
                return key_lower
            else:
                return self.search_lower(key,key_lower,avl[2])
        #avl[1] == keyの場合
        if avl[0] == None:
            return key_lower
        else:
            if key_lower == None:
                return self.end_higher(avl[0][1],avl[0])
            else:
                return max(key_lower,self.end_higher(avl[0][1],avl[0]))
    
    # keyより大きい最小の要素を検索する。なければNoneを返す。
    # 使用法 Instance.search_higher(x)
    def search_higher(self,key,key_higher=None,avl=0):
        if avl == 0:
            avl = self.avl
        if avl == None:
            return key_higher
        if  key < avl[1]:
            key_higher = avl[1]
            if avl[0] == None:
                return key_higher
            else:
                return self.search_higher(key,key_higher,avl[0])
        if avl[1] < key:
            if avl[2] == None:
                return key_higher
            else:
                return self.search_higher(key,key_higher,avl[2])
        #self.key == keyの場合
        if avl[2] == None:
            return key_higher
        else:
            if key_higher == None:
                return self.end_lower(avl[2][1],avl[2])
            else:
                return min(key_higher,self.end_lower(avl[2][1],avl[2]))
    
    def end_lower(self,end_lower_key,avl=None):
        if avl == None:
            avl = self.avl
        if avl[0] == None:
            return end_lower_key
        else:
            end_lower_key = avl[0][1]
            return self.end_lower(end_lower_key,avl[0])
    def end_higher(self,end_higher_key,avl=None):
        if avl == None:
            avl = self.avl
        if avl[2] == None:
            return end_higher_key
        else:
            end_higher_key = avl[2][1]
            return self.end_higher(end_higher_key,avl[2])

    def DoubleRightRotation(self,avl):
        tl = avl[0]
        avl[0] = tl[2][2]
        tl[2][2] = avl # selfはそのノード
        tlr = tl[2]
        tl[2] = tlr[0]
        tlr[0] = tl
        if tlr[3] == 1:
            tlr[2][3] = 2
            tlr[0][3] = 0
        elif tlr[3] == 2:
            tlr[2][3] = 0
            tlr[0][3] = 1
        elif tlr[3] == 0:
            tlr[2][3] = 0
            tlr[0][3] = 0
        tlr[3] = 0
        return tlr

    def DoubleLeftRotation(self,avl):
        tr = avl[2]
        avl[2] = tr[0][0]
        tr[0][0] = avl
        trl = tr[0]
        tr[0] = trl[2]
        trl[2] = tr
        if trl[3] == 2:
            trl[0][3] = 1
            trl[2][3] = 0
        elif trl.balance == 1:
            trl[0][3] = 0
            trl[2][3] = 2
        elif trl.balance == 0:
            trl[0][3] = 0
            trl[2][3] = 0
        trl[3] = 0
        return trl

    def SingleLeftRotation(self,avl):
        tr = avl[2]
        tr[3] = 0
        avl[3] = 0
        avl[2] = tr[0]
        tr[0] = avl
        return tr

    def SingleRightRotation(self,avl):
        tl = avl[0]
        tl[3] = 0
        avl[3] = 0
        avl[0] = tl[2]
        tl[2] = avl
        return tl

    def replace(self,p,v,avl): # 親ノードpの下にある自分(avl)をvに置き換える。
        if p[0] is avl:
            p[0] = v
        else :
            p[2] = v
    
    # 木に要素を追加する。
    # 使用法 Instance.insert(x)
    def insert(self,key): # rootでのみ呼ばれる挿入
        if self.avl[1] == None: # rootを含むrotationはしないことにする。
            self.avl[1] = key
            return self.avl
        if key < self.avl[1]:
            if self.avl[0] == None:
                self.avl[0] = self.node(key)
            else:
                self.insertx(self.avl,key,self.avl[0])
        elif self.avl[1] < key:
            if self.avl[2] == None:
                self.avl[2] = self.node(key)
            else:
                self.insertx(self.avl,key,self.avl[2])
        else: # key == self.avl[1]:
            pass # do not overwrite
    def insertx(self,p,key,avl): # replaceを呼ぶために一つ上の親を持っているinsert
        if key < avl[1]:
            if avl[0] == None:
                avl[0] = self.node(key)
            else:
                if not self.insertx(avl,key,avl[0]): # 左の木が生長しなければ、
                    return False # 成長しない
            balance = avl[3]
            if balance == 2:
                avl[3] = 0
                return False
            elif balance == 0:
                avl[3] = 1
                return True # 成長した
            elif balance == 1:
                if avl[0][3] == 2:
                    self.replace(p,self.DoubleRightRotation(avl),avl)
               
                return False # rotationを行うと成長しない
        if avl[1] < key:
            if avl[2] == None:
                avl[2] = self.node(key)
            else:
                if not self.insertx(avl,key,avl[2]):
                    return False
            balance = avl[3]
            if balance == 1:
                avl[3] = 0
                return False
            elif balance == 0:
                avl[3] = 2
                return True
            elif balance == 2:
                if avl[2][3] == 1:
                    self.replace(p,self.DoubleLeftRotation(avl),avl)
                elif avl[2][3] == 2:
                    self.replace(p,self.SingleLeftRotation(avl),avl)
                return False
        return False # avl[1] == keyの時は何もしないので成長もしない

    def debug(self):
        print(self.avl)
def main():
    n = int(input())
    p = list(map(int,input().split()))
    q = [0]*(n+1)
    for i in range(n):
        q[p[i]] = i+1
    
    ans = 0
    avl = Avltree()
    avl.insert(q[n])
    for i in range(n-1,0,-1):
        cnt = 0
        v = q[i]
        lo = avl.search_lower(v)
        hi = avl.search_higher(v)
        #print(lo,hi)
        if lo:
            lo2 = avl.search_lower(lo,0)
            t = lo-lo2
            if hi:
                t *= hi - v
            else:
                t *= n+1 - v
            cnt += t
        if hi:
            hi2 = avl.search_higher(hi,n+1)
            t = hi2 - hi
            if lo:
                t *= v - lo
            else:
                t *= v
            cnt += t
        #print(cnt)
        ans += i*cnt
        avl.insert(v)
    print(ans)

if __name__ == "__main__":
    main()

Buggy submission (verdict: Runtime Error):
class Avltree:
    """ Avl木を表現するクラス """
    def __init__(self):
        self.avl = [None,None,None,0]

    def node(self,key):
        return [None,key,None,0]

    # search(x) xが存在するならTrueをしないならFalseを返す
    # 使用法 Instance.search(x)
    def search(self,key,avl=None):
        if avl == None:
            avl = self.avl
        if key == avl[1]:
            return True
        elif key < avl[1]:
            if avl[0] == None:
                return False
            else:
                return self.search(key,avl[0])
        elif avl[1] < key:
            if avl[2] == None:
                return False
            else:
                return self.search(key,avl[2])
    
    # key未満で最大の要素を検索する。なければNoneを返す。
    # 使用法 Instance.search_lower(x)
    def search_lower(self,key,key_lower=None,avl=0):
        if avl == 0:
            avl = self.avl
        if avl == None:
            return key_lower
        elif key < avl[1]:
            if avl[0] == None:
                return key_lower
            else:
                return self.search_lower(key,key_lower,avl[0])
        elif avl[1] < key:
            key_lower = avl[1]
            if avl[2] == None:
                return key_lower
            else:
                return self.search_lower(key,key_lower,avl[2])
        #avl[1] == keyの場合
        if avl[0] == None:
            return key_lower
        else:
            if key_lower == None:
                return self.end_higher(avl[0][1],avl[0])
            else:
                return max(key_lower,self.end_higher(avl[0][1],avl[0]))
    
    # keyより大きい最小の要素を検索する。なければNoneを返す。
    # 使用法 Instance.search_higher(x)
    def search_higher(self,key,key_higher=None,avl=0):
        if avl == 0:
            avl = self.avl
        if avl == None:
            return key_higher
        if  key < avl[1]:
            key_higher = avl[1]
            if avl[0] == None:
                return key_higher
            else:
                return self.search_higher(key,key_higher,avl[0])
        if avl[1] < key:
            if avl[2] == None:
                return key_higher
            else:
                return self.search_higher(key,key_higher,avl[2])
        #self.key == keyの場合
        if avl[2] == None:
            return key_higher
        else:
            if key_higher == None:
                return self.end_lower(avl[2][1],avl[2])
            else:
                return min(key_higher,self.end_lower(avl[2][1],avl[2]))
    
    def end_lower(self,end_lower_key,avl=None):
        if avl == None:
            avl = self.avl
        if avl[0] == None:
            return end_lower_key
        else:
            end_lower_key = avl[0][1]
            return self.end_lower(end_lower_key,avl[0])
    def end_higher(self,end_higher_key,avl=None):
        if avl == None:
            avl = self.avl
        if avl[2] == None:
            return end_higher_key
        else:
            end_higher_key = avl[2][1]
            return self.end_higher(end_higher_key,avl[2])

    def DoubleRightRotation(self,avl):
        tl = avl[0]
        avl[0] = tl[2][2]
        tl[2][2] = avl # selfはそのノード
        tlr = tl[2]
        tl[2] = tlr[0]
        tlr[0] = tl
        if tlr[3] == 1:
            tlr[2][3] = 2
            tlr[0][3] = 0
        elif tlr[3] == 2:
            tlr[2][3] = 0
            tlr[0][3] = 1
        elif tlr[3] == 0:
            tlr[2][3] = 0
            tlr[0][3] = 0
        tlr[3] = 0
        return tlr

    def DoubleLeftRotation(self,avl):
        tr = avl[2]
        avl[2] = tr[0][0]
        tr[0][0] = avl
        trl = tr[0]
        tr[0] = trl[2]
        trl[2] = tr
        if trl[3] == 2:
            trl[0][3] = 1
            trl[2][3] = 0
        elif trl.balance == 1:
            trl[0][3] = 0
            trl[2][3] = 2
        elif trl.balance == 0:
            trl[0][3] = 0
            trl[2][3] = 0
        trl[3] = 0
        return trl

    def SingleLeftRotation(self,avl):
        tr = avl[2]
        tr[3] = 0
        avl[3] = 0
        avl[2] = tr[0]
        tr[0] = avl
        return tr

    def SingleRightRotation(self,avl):
        tl = avl[0]
        tl[3] = 0
        avl[3] = 0
        avl[0] = tl[2]
        tl[2] = avl
        return tl

    def replace(self,p,v,avl): # 親ノードpの下にある自分(avl)をvに置き換える。
        if p[0] is avl:
            p[0] = v
        else :
            p[2] = v
    
    # 木に要素を追加する。
    # 使用法 Instance.insert(x)
    def insert(self,key): # rootでのみ呼ばれる挿入
        if self.avl[1] == None: # rootを含むrotationはしないことにする。
            self.avl[1] = key
            return self.avl
        if key < self.avl[1]:
            if self.avl[0] == None:
                self.avl[0] = self.node(key)
            else:
                self.insertx(self.avl,key,self.avl[0])
        elif self.avl[1] < key:
            if self.avl[2] == None:
                self.avl[2] = self.node(key)
            else:
                self.insertx(self.avl,key,self.avl[2])
        else: # key == self.avl[1]:
            pass # do not overwrite
    def insertx(self,p,key,avl): # replaceを呼ぶために一つ上の親を持っているinsert
        if key < avl[1]:
            if avl[0] == None:
                avl[0] = self.node(key)
            else:
                if not self.insertx(avl,key,avl[0]): # 左の木が生長しなければ、
                    return False # 成長しない
            balance = avl[3]
            if balance == 2:
                avl[3] = 0
                return False
            elif balance == 0:
                avl[3] = 1
                return True # 成長した
            elif balance == 1:
                if avl[0][3] == 2:
                    self.replace(p,self.DoubleRightRotation(avl),avl)
               
                return False # rotationを行うと成長しない
        if avl[1] < key:
            if avl[2] == None:
                avl[2] = self.node(key)
            else:
                if not self.insertx(avl,key,avl[2]):
                    return False
            balance = avl[3]
            if balance == 1:
                avl[3] = 0
                return False
            elif balance == 0:
                avl[3] = 2
                return True
            elif balance == 2:
                if avl[2][3] == 1:
                    self.replace(p,self.DoubleLeftRotation(avl),avl)
                elif avl[2][3] == 2:
                    self.replace(p,self.SingleLeftRotation(avl),avl)
                return False
        return False # avl[1] == keyの時は何もしないので成長もしない

    def debug(self):
        print(self.avl)
def main():
    n = int(input())
    p = list(map(int,input().split()))
    q = [0]*(n+1)
    for i in range(n):
        q[p[i]] = i+1
    
    ans = 0
    avl = Avltree()
    avl.insert(q[n])
    for i in range(n-1,0,-1):
        cnt = 0
        v = q[i]
        lo = avl.search_lower(v)
        hi = avl.search_higher(v)
        #print(lo,hi)
        if lo:
            lo2 = avl.search_lower(lo,0)
            t = lo-lo2
            if hi:
                t *= hi - v
            else:
                t *= n+1 - v
            cnt += t
        if hi:
            hi2 = avl.search_higher(hi,n+1)
            t = hi2 - hi
            if lo:
                t *= v - lo
            else:
                t *= v
            cnt += t
        #print(cnt)
        ans += i*cnt
        avl.insert(v)
    print(ans)

if __name__ == "__main__":
    main()