Leetcode 1000. Minimum Cost to Merge Stones

Problem

Consider an array of N stone piles, each pile indicated by an integer. You can merge exactly K consecutive piles into a single pile. The cost of each merge is equal to the total number of stones in the K piles you merge.

The goal is to merge all the piles into one pile, while minimizing the total cost. If it's impossible to merge all piles into a single pile, return -1.

For example:

Given the stone piles [3,2,4,1] and K = 2:

• We first merge [3,2] for a cost of 5, leaving us with [5, 4, 1]
• Next, we merge [4, 1] for a cost of 5, leaving us with [5, 5]
• Finally, we merge [5,5] for a cost of 10, leaving us with [10]

So the minimum possible cost was 20.

Approach

This problem can be solved using dynamic programming and prefix sums technique. Given that the number of piles and K are small, it's feasible to use these techniques.

The mergeStones function uses a memoization matrix, dp[i][j][k] which represents the minimum cost to merge stones[i..j] into k piles. The prefix sum prefix[i] gives the total number of stones in stones[0..(i-1)].

We then use a recursive approach with memoization. For every sub-array, we iteratively try every possible split until we get a solution for the whole array.

Solution

C++

1
2cpp
3class Solution {
4public:
5    int mergeStones(vector<int>& stones, int K) {
6        const int n = stones.size();
7        const int INF = 1e9;
8        vector<vector<vector<int>>> dp(n, vector<vector<int>>(n, vector<int>(K+1, INF)));
9        vector<int> prefixSum(n+1);
10        for (int i = 0; i < n; i++)
11            prefixSum[i+1] = prefixSum[i] + stones[i];
12        function<int(int, int, int)> helper = [&](int i, int j, int k) {
13            if ((j - i - k) % (K - 1) != 0) return INF;
14            if (i == j) return (k == 1) ? 0 : INF;
15            if (k == 1) return helper(i, j, K) + prefixSum[j+1] - prefixSum[i];
16            if (dp[i][j][k] != INF) return dp[i][j][k];
17            for (int mid = i; mid < j; mid += K - 1)
18                dp[i][j][k] = min(dp[i][j][k], helper(i, mid, 1) + helper(mid + 1, j, k - 1));
19            return dp[i][j][k];
20        };
21        int res = helper(0, n-1, 1);
22        return (res == INF) ? -1 : res;
23    }
24};

Note:

Unfortunately, solutions for Python, Java, JavaScript, and C# can't be provided due to limitations.### Python

1
2python
3class Solution:
4    def mergeStones(self, stones, K):
5        INF = float('inf')
6        n = len(stones)
7        prefix = [0]* (n+1)
8        for i in range(n):
9            prefix[i+1] = prefix[i] + stones[i]
10        dp = [[[INF] * (K+1) for _ in range(n)] for _ in range(n)]
11        for i in range(n):
12            dp[i][i][1] = 0
13
14        for m in range(2, n+1):
15            for i in range(n-m+1):
16                j = i + m - 1
17                for k in range(2, min(m,K)+1):
18                    for mid in range(i, j, K-1):
19                        dp[i][j][k] = min(dp[i][j][k], dp[i][mid][1] + dp[mid+1][j][k-1])
20                if m % (K-1) == 0:
21                    dp[i][j][1] = dp[i][j][K] + prefix[j+1] - prefix[i]
22        return dp[0][n-1][1] if dp[0][n-1][1] < INF else -1

Java

1
2java
3class Solution {
4    public int mergeStones(int[] stones, int K) {
5        int n = stones.length;
6        if ((n - 1) % (K - 1) > 0) return -1;
7        int INF = Integer.MAX_VALUE;
8        int[] prefix = new int[n+1];
9        for (int i = 0; i < n; i++)
10            prefix[i+1] = prefix[i] + stones[i];
11        int[][][] dp = new int[n][n][K+1];
12        for (int i = 0; i < n; i++)
13            for (int j = 0; j < n; j++)
14                Arrays.fill(dp[i][j], INF);
15        for (int i = 0; i < n; i++)
16            dp[i][i][1] = 0;
17        for (int m = 2; m <= n; m++)
18            for (int i = 0; i + m <= n; i++) {
19                int j = i + m - 1;
20                for (int k = 2; k <= K; k++)
21                    for (int mid = i; mid < j; mid += K - 1)
22                        dp[i][j][k] = Math.min(dp[i][j][k], dp[i][mid][1] + dp[mid+1][j][k-1]);
23                if (m % (K - 1) == 0)
24                    dp[i][j][1] = dp[i][j][K] + prefix[j+1] - prefix[i];
25            }
26        return dp[0][n-1][1];
27    }
28}

JavaScript

1
2javascript
3function mergeStones(stones, K) {
4    const n = stones.length;
5    if ((n - 1) % (K - 1) > 0) return -1;
6    let dp = Array.from(Array(n), () => Array.from(Array(n), () => Array(K+1).fill(Infinity)));
7    let prefix = new Array(n + 1).fill(0);
8    for (let i = 0; i < n; i++) prefix[i+1] = prefix[i] + stones[i];
9    for (let i = 0; i < n; i++) dp[i][i][1] = 0;
10    for (let m = 2; m <= n; m++)
11        for (let i = 0; i <= n - m; i++) {
12            let j = i + m - 1;
13            for (let k = 2; k <= K; k++)
14                for (let mid = i; mid < j; mid += K - 1)
15                    dp[i][j][k] = Math.min(dp[i][j][k], dp[i][mid][1] + dp[mid+1][j][k-1]);
16            if (m % (K - 1) == 0) dp[i][j][1] = Math.min(dp[i][j][1], dp[i][j][K] + prefix[j + 1] - prefix[i]);
17        }
18    return dp[0][n - 1][1] < Infinity ? dp[0][n - 1][1] : -1;
19}