Teaching Kids Programming – Minimum Cost to Connect Sticks (Priority Queue/Min Heap)

Teaching Kids Programming: Videos on Data Structures and Algorithms

You have some number of sticks with positive integer lengths. These lengths are given as an array sticks, where sticks[i] is the length of the ith stick. You can connect any two sticks of lengths x and y into one stick by paying a cost of x + y. You must connect all the sticks until there is only one stick remaining.

Return the minimum cost of connecting all the given sticks into one stick in this way.

Example 1:
Input: sticks = [2,4,3]
Output: 14
Explanation: You start with sticks = [2,4,3].
1. Combine sticks 2 and 3 for a cost of 2 + 3 = 5. Now you have sticks = [5,4].
2. Combine sticks 5 and 4 for a cost of 5 + 4 = 9. Now you have sticks = [9].
There is only one stick left, so you are done. The total cost is 5 + 9 = 14.

Example 2:
Input: sticks = [1,8,3,5]
Output: 30
Explanation: You start with sticks = [1,8,3,5].
1. Combine sticks 1 and 3 for a cost of 1 + 3 = 4. Now you have sticks = [4,8,5].
2. Combine sticks 4 and 5 for a cost of 4 + 5 = 9. Now you have sticks = [9,8].
3. Combine sticks 9 and 8 for a cost of 9 + 8 = 17. Now you have sticks = [17].
There is only one stick left, so you are done. The total cost is 4 + 9 + 17 = 30.

Example 3:
Input: sticks = [5]
Output: 0
Explanation: There is only one stick, so you don’t need to do anything. The total cost is 0.

Constraints:
1 <= sticks.length <= 104
1 <= sticks[i] <= 104

Hints:
How many times does every stick contribute to the answer?
Some of the sticks will be used more than the others. Which sticks should be used the most/least?
The sticks with long lengths cost a lot so we should use these the least.
What if we keep merging the two shortest sticks?

Greedy Strategy by Picking Two Smallest

Assuming we are given four sticks with lengths a, b, c, and d respectively. The total costs would be (merging from left to right):

connect a and b with cost a + b;
connect a+b and c with cost a + b + c;
connect a + b + c and d with cost a + b + c + d;

Total costs: 3a + 3b + 2c + d. Thus, it is optimial to greedily pick the two smallest (a and b) every round in order to achieve the minimal total costs.

By using the min heap in Python (the heapify module), we can simulate the process by popping out two smallest from the Priority Queue, accumulate the cost, and add the new stick (merged length) back to the priority queue.

1
2
3
4
5
6
7
8
9
10
11

import heapq as hq
 
class Solution:
    def connectSticks(self, sticks: List[int]) -> int:
        hq.heapify(sticks)        
        ans = 0
        while len(sticks) > 1:
            p, q = hq.heappop(sticks), hq.heappop(sticks)
            ans += p + q
            hq.heappush(sticks, p + q)
        return ans

import heapq as hq

class Solution:
    def connectSticks(self, sticks: List[int]) -> int:
        hq.heapify(sticks)        
        ans = 0
        while len(sticks) > 1:
            p, q = hq.heappop(sticks), hq.heappop(sticks)
            ans += p + q
            hq.heappush(sticks, p + q)
        return ans

The time complexity is O(NLogN) where N is the number of the sticks, and O(LogN) for the costs of pushing/popping element to/from the heap. Space complexity is O(N) for the heap/priority queue.

Related Post: Compute the Minimum Costs to Connect the Sticks using Priority Queue or Heap

–EOF (The Ultimate Computing & Technology Blog)