牛客网 构造maxTree

牛客网 构造maxTree

题目描述
对于一个没有重复元素的整数数组，请用其中元素构造一棵MaxTree，MaxTree定义为一棵二叉树，其中的节点与数组元素一一对应，同时对于MaxTree的每棵子树，它的根的元素值为子树的最大值。现有一建树方法，对于数组中的每个元素，其在树中的父亲为数组中它左边比它大的第一个数和右边比它大的第一个数中更小的一个。若两边都不存在比它大的数，那么它就是树根。请证明这个方法的正确性，同时设计O(n)的算法实现这个方法。
给定一个无重复元素的数组A和它的大小n，请返回一个数组，其中每个元素为原数组中对应位置元素在树中的父亲节点的编号，若为根则值为-1。
测试样例：
[3,1,4,2],4
返回：[2,0,-1,2]

分析：先利用栈两次循环分别求出左边第一个大于i的数组lt[]和右边第一个大于i的数组rt[]。最后再一次循环比较两个数组得出结果
代码如下：

class MaxTree:
    def buildMaxTree(self, A, n):
        # write code here
        if n <= 0:
            retturn []
        ret = [-1] * n
        stack = []
        lt = [-1]*n; rt = [-1]*n
        for i in range(n):
            #l,r = i - 1, i + 1
            while(len(stack) > 0 and A[i] > A[stack[-1]]):
                stack.pop()
            if len(stack) > 0:
                lt[i] = stack[-1]
            stack.append(i)

        while len(stack)> 0:
            stack.pop()

        for i in range(n-1, -1, -1):
            while(len(stack) > 0 and A[i] > A[stack[-1]]):
                stack.pop()
            if len(stack) > 0:
                rt[i] = stack[-1]
            stack.append(i)
        for i in range(n):
            if(lt[i] != -1 and rt[i] != -1):
                if(A[lt[i]] < A[rt[i]]):
                    ret[i] = lt[i]
                else:
                    ret[i] = rt[i]
            elif(lt[i] != -1):
                ret[i] = lt[i]
            elif(rt[i] != -1):
                ret[i] = rt[i]
            else:
                pass
        return ret