Problem statement

https://leetcode.com/problems/minimum-cost-tree-from-leaf-values/

Solution 1

First solution is using dp with states dp(i, j) is answer for arr[i: j + 1].

Complexity

It is O(n^3) for time and O(n^2) for space.

Code

class Solution:
    def mctFromLeafValues(self, arr):
        n = len(arr)
        m = {(i, j): max(arr[i:j+1]) for i in range(n) for j in range(i, n)}
        
        @lru_cache(None)
        def dp(i, j):
            if i == j: return arr[i]
            return min(dp(i, k) + dp(k + 1, j) + m[i, k]*m[k+1, j] for k in range(i, j))
            
        return dp(0, n - 1) - sum(arr)

Solution 2

In fact we can do better, using monotonic stack! The idea is that for each element arr[i] we need to find the closest element to the left and closest element to the right which is >= arr[i]. Then the answer will be sum of arr[i] * min(lft[i], rgh[i]). However it is a good question, why, it is greedy idea, proof TBD. For this we can reuse problem Leetcode 0907 Sum of Subarray Minimums.

Complexity

It is O(n) for time and space.

Code

class Solution:
    def mctFromLeafValues(self, arr):
        n = len(arr)
        rgh, lft = [float("inf")]*n, [float("inf")]*n

        stack1 = []
        for i in range(n):
            while stack1 and arr[i] >= arr[stack1[-1]]: 
                idx = stack1.pop()
                rgh[idx] = arr[i]
            stack1.append(i)
            
        stack2 = []
        for i in range(n-1, -1, -1):
            while stack2 and arr[i] > arr[stack2[-1]]:
                idx = stack2.pop()
                lft[idx] = arr[i]
            stack2.append(i)

        ans = 0
        for x, y, a in zip(lft, rgh, arr):
            if min(x, y) != float("inf"): ans += min(x, y) * a
                
        return ans