求解最小化堆合并成本的算法问题

Got it, let's break down this problem and figure out how to solve it efficiently. First, let's restate the problem clearly to make sure we're aligned:

We have N piles of items, where the i-th pile has M[i] items. We can only merge pile i into pile j if j > i, and the cost of this merge is M[i]*(j - i). Our goal is to merge the initial N piles into exactly K piles with the minimum total cost.

Key Insight #

Here's a crucial observation that simplifies everything: merging pile i into j directly has the same total cost as merging it through intermediate piles. For example, merging i into k then k into j costs M[i](k-i) + M[i](j-k) = M[i]*(j-i) — exactly the same as merging i straight into j.

This means to minimize cost, every non-kept pile should be merged into the closest kept pile to its right (since a smaller distance means lower cost). So the problem reduces to selecting K positions in the original sequence (as our kept piles, with the last pile always being kept—since we can't merge it into anything else) such that the sum of M[i]*(closest_right_kept_pos - i) for all non-kept i is minimized.

Dynamic Programming Approach #

We'll use dynamic programming (DP) to model this problem:

State Definition #

Let dp[i][j] represent the minimum cost to keep exactly j piles among the first i piles, where the j-th kept pile is the i-th pile. Our target answer is dp[N][K] (since the last pile must be kept).

Transition Equation #

To compute dp[i][j], we need to consider all possible positions k (where j-1 ≤ k < i) where we kept j-1 piles in the first k piles. The cost added by choosing k as the previous kept pile is the total cost of merging all piles from k+1 to i-1 into i.

To calculate this cost quickly, we precompute two prefix sums:

• S[i]: Sum of M[1] to M[i] (1-based indexing for easier calculations)
• T[i]: Sum of M[x]*x for x from 1 to i

The cost of merging piles k+1 to i-1 into i is:
i*(S[i-1] - S[k]) - (T[i-1] - T[k])

So the transition becomes:

dp[i][j] = min{ dp[k][j-1] + i*(S[i-1] - S[k]) - (T[i-1] - T[k]) } for all k in [j-1, i-1]

Initialization #

• dp[i][1]: Cost to keep only the i-th pile from the first i piles. This is the sum of merging all piles 1 to i-1 into i, which simplifies to i*S[i-1] - T[i-1].
• dp[1][1] = 0: No cost needed if we only have one pile.

Implementation (Basic Version) #

Here's a Python implementation of the basic DP approach (time complexity: O(N²K)):

def main():
    import sys
    input = sys.stdin.read().split()
    idx = 0
    n = int(input[idx])
    idx += 1
    k = int(input[idx])
    idx += 1
    m = list(map(int, input[idx:idx+n]))
    
    # Convert to 1-based indexing
    m = [0] + m
    S = [0] * (n + 1)
    T = [0] * (n + 1)
    
    for i in range(1, n+1):
        S[i] = S[i-1] + m[i]
        T[i] = T[i-1] + m[i] * i
    
    INF = float('inf')
    dp = [[INF] * (k + 1) for _ in range(n + 1)]
    
    # Initialize for j=1
    for i in range(1, n+1):
        dp[i][1] = i * S[i-1] - T[i-1]
    dp[1][1] = 0  # Edge case: single pile
    
    # Fill dp table for j >=2
    for j in range(2, k+1):
        for i in range(j, n+1):  # Need at least j piles to keep j piles
            for k_prev in range(j-1, i):
                if dp[k_prev][j-1] == INF:
                    continue
                cost = i * (S[i-1] - S[k_prev]) - (T[i-1] - T[k_prev])
                if dp[k_prev][j-1] + cost < dp[i][j]:
                    dp[i][j] = dp[k_prev][j-1] + cost
    
    print(dp[n][k])

if __name__ == "__main__":
    main()

Optimization with Quadrangle Inequality #

For larger N (e.g., N=1e3 or more), the O(N²K) approach might be too slow. We can optimize this using the quadrangle inequality, which tells us that the optimal k for dp[i][j] (let's call it opt[i][j]) satisfies opt[i-1][j] ≤ opt[i][j] ≤ opt[i][j+1]. This reduces the number of k values we need to check per state, bringing the time complexity down to O(NK).

Here's the optimized implementation:

def main():
    import sys
    input = sys.stdin.read().split()
    idx = 0
    n = int(input[idx])
    idx += 1
    k = int(input[idx])
    idx += 1
    m = list(map(int, input[idx:idx+n]))
    
    # Convert to 1-based indexing
    m = [0] + m
    S = [0] * (n + 1)
    T = [0] * (n + 1)
    
    for i in range(1, n+1):
        S[i] = S[i-1] + m[i]
        T[i] = T[i-1] + m[i] * i
    
    INF = float('inf')
    dp = [[INF] * (k + 1) for _ in range(n + 2)]  # Extra space for boundary
    opt = [[0] * (k + 1) for _ in range(n + 2)]   # Stores optimal k_prev for each state
    
    dp[0][0] = 0  # Base case: 0 piles, 0 kept, cost 0
    
    # Initialize for j=1
    for i in range(1, n+1):
        dp[i][1] = i * S[i-1] - T[i-1]
        opt[i][1] = 0
    
    # Fill dp table for j >=2
    for j in range(2, k+1):
        opt[n+1][j] = n  # Boundary condition
        # Iterate i from n down to j to use previous opt values
        for i in range(n, j-1, -1):
            start = opt[i-1][j]
            end = opt[i+1][j]
            min_cost = INF
            best_k = start
            # Only check k_prev in [start, end]
            for k_prev in range(start, end+1):
                if dp[k_prev][j-1] == INF:
                    continue
                current_cost = dp[k_prev][j-1] + i*(S[i-1]-S[k_prev]) - (T[i-1]-T[k_prev])
                if current_cost < min_cost:
                    min_cost = current_cost
                    best_k = k_prev
            dp[i][j] = min_cost
            opt[i][j] = best_k
    
    print(dp[n][k])

if __name__ == "__main__":
    main()

Final Notes #

• The key insight about merge costs being equivalent regardless of intermediate steps is what makes this problem tractable.
• Prefix sums are essential to avoid recalculating sums repeatedly, which would otherwise make the solution too slow.
• The quadrangle inequality optimization is a powerful tool for reducing the time complexity of DP problems with monotonic optimal decisions.

内容的提问来源于stack exchange，提问作者Jingjie Yang