1. 程式人生 > >構造MaxTree(牛客網)

構造MaxTree(牛客網)

題目描述

對於一個沒有重複元素的整數陣列,請用其中元素構造一棵MaxTree,MaxTree定義為一棵二叉樹,其中的節點與陣列元素一一對應,同時對於MaxTree的每棵子樹,它的根的元素值為子樹的最大值。現有一建樹方法,對於陣列中的每個元素,其在樹中的父親為陣列中它左邊比它大的第一個數和右邊比它大的第一個數中更小的一個。若兩邊都不存在比它大的數,那麼它就是樹根。請證明這個方法的正確性,同時設計O(n)的演算法實現這個方法。

給定一個無重複元素的陣列A和它的大小n,請返回一個數組,其中每個元素為原陣列中對應位置元素在樹中的父親節點的編號,若為根則值為-1。

測試樣例:

[3,1,4,2],4
返回:[2,0,-1,2]

證明過程詳見左程雲的《程式設計師程式碼面試指南》P22-24

下面給出自己改寫的程式碼

import java.util.*;

public class MaxTree {
    public int[] buildMaxTree(int[] A, int n) {
        // write code here
        int[] res = new int[n];
		Stack<Integer> stack = new Stack<>();

		for (int i = 0; i < n; i++) {
			while (!stack.isEmpty() && A[i] > A[stack.peek()]) {
				int index = stack.pop();
				if (stack.isEmpty()) {
					res[index] = i;
				} else {
					res[index] = A[stack.peek()] < A[i] ? stack.peek() : i;
				}
			}
			stack.push(i);
		}

		while (!stack.isEmpty()) {
			int index = stack.pop();
			if (stack.isEmpty()) {
				res[index] = -1;
			} else {
				res[index] = stack.peek();
			}
		}

		return res;
    }
}