Problem p03987 (python)

Fixed submission:
# -*- coding: utf-8 -*-
import sys
# sys.setrecursionlimit(10**6)
# buff_readline = sys.stdin.buffer.readline
buff_readline = sys.stdin.readline
readline = sys.stdin.readline

INF = 2**62-1


def read_int():
    return int(buff_readline())


def read_int_n():
    return list(map(int, buff_readline().split()))


def read_float():
    return float(buff_readline())


def read_float_n():
    return list(map(float, buff_readline().split()))


def read_str():
    return readline().strip()


def read_str_n():
    return readline().strip().split()


def error_print(*args):
    print(*args, file=sys.stderr)


def mt(f):
    import time

    def wrap(*args, **kwargs):
        s = time.time()
        ret = f(*args, **kwargs)
        e = time.time()

        error_print(e - s, 'sec')
        return ret

    return wrap


class SegmentTree:
    def __init__(self, array, operator, identity_element):
        _len = len(array)
        self.__op = operator
        self.__size = 1 << (_len - 1).bit_length()
        self.__tree = [identity_element] * self.__size + \
            array + [identity_element] * (self.__size - _len)
        self.__ie = identity_element

        for i in range(self.__size - 1, 0, -1):
            self.__tree[i] = operator(
                self.__tree[i * 2], self.__tree[i * 2 + 1])

    def update(self, i, v):
        i += self.__size
        self.__tree[i] = v
        while i:
            i //= 2
            self.__tree[i] = self.__op(
                self.__tree[i * 2], self.__tree[i * 2 + 1])

    def query(self, l, r):
        """[l, r)
        """
        l += self.__size
        r += self.__size
        ret = self.__ie
        while l < r:
            if l & 1:
                ret = self.__op(ret, self.__tree[l])
                l += 1
            if r & 1:
                r -= 1
                ret = self.__op(ret, self.__tree[r])
            l //= 2
            r //= 2
        return ret


@mt
def slv(N, A):
    ai = [(a, i) for i, a in enumerate(A)]
    ai.sort()

    lst = SegmentTree([-1] * N, max, -1)
    rst = SegmentTree([N] * N, min, N)
    ans = 0
    for a, i in ai:
        l = lst.query(0, i)
        r = rst.query(i+1, N)
        ans += (i-l) * (r-i) * a
        lst.update(i, i)
        rst.update(i, i)
    return ans


def main():
    N = read_int()
    A = read_int_n()
    print(slv(N, A))


if __name__ == '__main__':
    main()

Buggy submission (verdict: Runtime Error):
# -*- coding: utf-8 -*-
import bisect
import heapq
import math
import random
import sys
from collections import Counter, defaultdict, deque
from decimal import ROUND_CEILING, ROUND_HALF_UP, Decimal
from functools import lru_cache, reduce
from itertools import combinations, combinations_with_replacement, product, permutations
from operator import add, mul, sub

sys.setrecursionlimit(10000)


def read_int():
    return int(input())


def read_int_n():
    return list(map(int, input().split()))


def read_float():
    return float(input())


def read_float_n():
    return list(map(float, input().split()))


def read_str():
    return input().strip()


def read_str_n():
    return list(map(str, input().split()))


def error_print(*args):
    print(*args, file=sys.stderr)


def mt(f):
    import time

    def wrap(*args, **kwargs):
        s = time.time()
        ret = f(*args, **kwargs)
        e = time.time()

        error_print(e - s, 'sec')
        return ret

    return wrap



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
        # print(avl, key)
        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)
                elif avl[0][3] == 1:
                    self.replace(p,self.SingleRightRotation(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)

@mt
def slv(N, A):
    ans = 0
    a2i = [0] * (N+1)
    for i, a in enumerate(A):
        a2i[A[i]] = i

    at = Avltree()
    at.insert(-1)
    at.insert(N)
    for a in range(1, N+1):
        i = a2i[a]
        l = at.search_lower(i)
        r = at.search_higher(i)
        # print(l, r, i, a)
        ans += (i - l) * (r - i) * a
        at.insert(i)

    return ans


def main():
    N = read_int()
    A = read_int_n()
    print(slv(N, A))

    # A = list(range(1, 200000+1))
    # for _ in range(100):
    #     random.shuffle(A)
    #     N = len(A)
    #     print(slv(N, A))


if __name__ == '__main__':
    main()

Buggy submission (verdict: Runtime Error):
# -*- coding: utf-8 -*-
import bisect
import heapq
import math
import random
import sys
from collections import Counter, defaultdict, deque
from decimal import ROUND_CEILING, ROUND_HALF_UP, Decimal
from functools import lru_cache, reduce
from itertools import combinations, combinations_with_replacement, product, permutations
from operator import add, mul, sub

sys.setrecursionlimit(10000)


def read_int():
    return int(input())


def read_int_n():
    return list(map(int, input().split()))


def read_float():
    return float(input())


def read_float_n():
    return list(map(float, input().split()))


def read_str():
    return input().strip()


def read_str_n():
    return list(map(str, input().split()))


def error_print(*args):
    print(*args, file=sys.stderr)


def mt(f):
    import time

    def wrap(*args, **kwargs):
        s = time.time()
        ret = f(*args, **kwargs)
        e = time.time()

        error_print(e - s, 'sec')
        return ret

    return wrap


@mt
def slv(N, A):
    import numpy as np
    A = np.array(A)

    def f(A):
        n = len(A)
        if not n:
            return 0
        if n == 1:
            return A[0]

        i = np.argmin(A)
        return A[i]*(i+1)*(n-i) + f(A[:i]) + f(A[i+1:])

    return f(A)


def main():
    N = read_int()
    A = read_int_n()
    print(slv(N, A))


if __name__ == '__main__':
    main()