Minimum Cost Tree From Leaf Values LeetCode solution in Java

1
class Solution {
2
  // Use stack. Similar to trapping the rain water problem and the largest rectangle under histogram
3
  // Use stack to keep a decreasing order by adding smaller values, while there is bigger value
4
  // arr[i] than the peek, pop it and store as mid and calculate the multiplication mid*min(arr[i],
5
  // stack.peek()).
6

7
  // NOTE: if we observe the number array, in order to obtain the smallest sum of all non-leaf
8
  // values, we want to merge those small values first. In order words, the smaller a value
9
  // is, the lower leaf it should stay because this way as we are building the tree up,
10
  // we are building smaller multiplication/parent node first as it is only going to get bigger
11
  // as we build the tree up.
12

13
  // Ex: 4 3 2 1 5
14
  // There are many ways we can build a tree following the problem's requirement. However, to
15
  // gain smallest sum. We need to merge 2 and 1 first as they are the two smallest ones. To
16
  // do that, we use the stack mentioned above as a decreasing order. After
17
  // that we get a parent node with value 2. This node could be a left or right child of its parent
18
  // but what we want is that its parent needs also be as small as possible. We also know that its
19
  // parent has one mutiplier already: 2. Note: this 2 is not from the product of 1 * 2, but from
20
  // the max  of child
21
  // 1 and 2 as the problem requires. So, we see what values next to the leaf 2 could be a
22
  // candidate. Obviously, 3 since it is the smallest one in the stack Then, 3
23
  // becomes the left child and 1*2 = 2 becomes right child. See below:
24
  //  ...
25
  //  / \
26
  // 3   2
27
  //    / \
28
  //   2   1
29
  //
30

31
  // If we observe carefully, 3 2 1 is decreasing... So how about every time we see a "dip" point
32
  // in the array we calculate its multiplication. To do that, say we are at arr[i] and their
33
  // relations are arr[i-1] <= arr[i] <= arr[i+1]. The min multiplication is a[i] * min(arr[i-1],
34
  // arr[i+1]). Then the example above is arr[i] = 1, arr[i-1] = 2, arr[i+1] = 5
35

36
  public int mctFromLeafValues(int[] arr) {
37
    if (arr == null || arr.length < 2) {
38
      return 0;
39
    }
40

41
    int res = 0;
42
    Stack<Integer> stack = new Stack<>();
43
    for (int num : arr) {
44

45
      // while num is bigger than peek(), pop and calculate
46
      while (!stack.isEmpty() && stack.peek() <= num) {
47
        int mid = stack.pop();
48
        if (stack.isEmpty()) res += mid * num;
49
        else res += mid * Math.min(stack.peek(), num);
50
      }
51

52
      stack.push(num); // if num is smaller, push into stack
53
    }
54

55
    // if there are values left in the stack, they sure will be mutiplied anyway
56
    // and added to the result.
57
    while (stack.size() > 1) { // > 1 because we have a peek() after pop() below
58
      res += stack.pop() * stack.peek();
59
    }
60

61
    return res;
62
  }
63
}