[LeetCode] 1000. Minimum Cost to Merge Stones

There are N piles of stones arranged in a row. The i-th pile has stones[i] stones.

A move consists of merging exactly K consecutive piles into one pile, and the cost of this move is equal to the total number of stones in these K piles.

Find the minimum cost to merge all piles of stones into one pile. If it is impossible, return -1.

Example 1:

Input: stones = [3,2,4,1], K = 2
Output: 20
Explanation: 
We start with [3, 2, 4, 1].
We merge [3, 2] for a cost of 5, and we are left with [5, 4, 1].
We merge [4, 1] for a cost of 5, and we are left with [5, 5].
We merge [5, 5] for a cost of 10, and we are left with [10].
The total cost was 20, and this is the minimum possible.

Example 2:

Input: stones = [3,2,4,1], K = 3
Output: -1
Explanation: After any merge operation, there are 2 piles left, and we can't merge anymore.  So the task is impossible.

Example 3:

Input: stones = [3,5,1,2,6], K = 3
Output: 25
Explanation: 
We start with [3, 5, 1, 2, 6].
We merge [5, 1, 2] for a cost of 8, and we are left with [3, 8, 6].
We merge [3, 8, 6] for a cost of 17, and we are left with [17].
The total cost was 25, and this is the minimum possible.

Note:

1 <= stones.length <= 30
2 <= K <= 30
1 <= stones[i] <= 100

class Solution {
    private Integer[][][] dp;
    private int[] prefixSum;
    public int mergeStones(int[] stones, int K) {
        int n = stones.length;
        dp = new Integer[n][n][K + 1];
        for(int i = 0; i < n; i++) {
            dp[i][i][1] = 0; 
        }
        prefixSum = new int[n + 1];
        for(int i = 1; i <= n; i++){
            prefixSum[i] = prefixSum[i - 1] + stones[i - 1];
        }
        int res = merge(stones, 0, n - 1, 1, K);
        return res == Integer.MAX_VALUE ? -1 : res;
    }
    private int merge(int[] stones, int start, int end, int numOfPiles, int K) {
        if((end - start + 1 - numOfPiles) % (K - 1) != 0) {
            return Integer.MAX_VALUE;
        }
        if(dp[start][end][numOfPiles] != null) {
            return dp[start][end][numOfPiles];
        }
        int minCost = Integer.MAX_VALUE;
        if(numOfPiles == 1) {
            int mergeCost = merge(stones, start, end, K, K);
            if(mergeCost != Integer.MAX_VALUE) {
                minCost = mergeCost + prefixSum[end + 1] - prefixSum[start];
            }
        }
        else {
            for(int mid = start; mid < end; mid++) {
                int leftMergeCost = merge(stones, start, mid, 1, K);
                if(leftMergeCost == Integer.MAX_VALUE) {
                    continue;
                } 
                int rightMergeCost = merge(stones, mid + 1, end, numOfPiles - 1, K);
                if(rightMergeCost == Integer.MAX_VALUE) {
                    continue;
                }
                minCost = Math.min(minCost, leftMergeCost + rightMergeCost);
            }            
        }
        dp[start][end][numOfPiles] = minCost;
        return minCost;
    }
}