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

You have some sticks with positive integer lengths. You can connect any two sticks of lengths X and Y into one stick by paying a cost of X + Y. You perform this action until there is 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

Example 2:
Input: sticks = [1,8,3,5]
Output: 30

Constraints:
1 <= sticks.length <= 10^4
1 <= sticks[i] <= 10^4

Using Priority Queue to Compute the Minimum Costs

Given sticks of length [2, 4, 3], the minimum costs would be (2+3)+(2+3+4) which is to connect (2,3) then connect to 4. The strategy would be to choose the 2 sticks that have two minimum lengths, then push the new stick back to the queue.

We can use a priority queue to do this. In C++, the priority queue by default pops out the current maximum, but we can easily specify a comparator e.g. std::greater that allows the priority queue to de-queue the current minimum for the queue.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16

class Solution {
public:
    int connectSticks(vector<int>& sticks) {
        priority_queue<int , vector<int>, greater<int> pq(begin(sticks), end(sticks)); 
        int res = 0; 
        while (pq.size() > 1) { // while still more than 1 stick
            int first = pq.top(); // get current minimum
            pq.pop(); 
            int second = pq.top(); // get current minimum
            pq.pop(); 
            res += first + second; // update the cost
            pq.push(first + second); // push the new stick back
        } 
        return res;        
    }
};

class Solution {
public:
    int connectSticks(vector<int>& sticks) {
        priority_queue<int , vector<int>, greater<int> pq(begin(sticks), end(sticks)); 
        int res = 0; 
        while (pq.size() > 1) { // while still more than 1 stick
            int first = pq.top(); // get current minimum
            pq.pop(); 
            int second = pq.top(); // get current minimum
            pq.pop(); 
            res += first + second; // update the cost
            pq.push(first + second); // push the new stick back
        } 
        return res;        
    }
};

Building a priority queue or heap takes O(N.LogN) complexity for N elements. Dequing the current minimum takes O(1) complexity and re-building (pushing a new stick back) takes O(logN), therefore the overall complexity is O(N.LogN).

The Heaps are binary trees that the parent nodes have values less or equal to its children. In Python, the min heap (priority queues) can be built using the heapq functions. For example, the heapq.heapify will transform a list into heap, in place. And the heapq.heappop returns the current minimum element in the heap while heapq.heappush adds a new item to the heap.

1
2
3
4
5
6
7
8
9
10

class Solution:
    def connectSticks(self, sticks: List[int]) -> int:
        heapq.heapify(sticks)
        res = 0
        while len(sticks) > 1:
            a = heapq.heappop(sticks)
            b = heapq.heappop(sticks)
            heapq.heappush(sticks, a + b)
            res += a + b
        return res

class Solution:
    def connectSticks(self, sticks: List[int]) -> int:
        heapq.heapify(sticks)
        res = 0
        while len(sticks) > 1:
            a = heapq.heappop(sticks)
            b = heapq.heappop(sticks)
            heapq.heappush(sticks, a + b)
            res += a + b
        return res

The following shows a max heap that root values are equal or greater than children nodes.

A Max Heap

Using Priority Queue in Java

We can use the PriorityQueue class in Java. The constructor takes two parameters: the first one is the capacity, and the second one is the optional comparator. We can pass in a lambda function asking for the reverse order.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16

class Solution {
    public int connectSticks(int[] sticks) {
        PriorityQueue<Integer> pq = new PriorityQueue<>(sticks.length, (a, b) -> a - b);
        for (int a: sticks) {
            pq.add(a);
        }
        int cost = 0;
        while (pq.size() > 1) {
            int a = pq.poll();
            int b = pq.poll();
            cost += a + b;
            pq.offer(a + b);
        }
        return cost;
    }
}

class Solution {
    public int connectSticks(int[] sticks) {
        PriorityQueue<Integer> pq = new PriorityQueue<>(sticks.length, (a, b) -> a - b);
        for (int a: sticks) {
            pq.add(a);
        }
        int cost = 0;
        while (pq.size() > 1) {
            int a = pq.poll();
            int b = pq.poll();
            cost += a + b;
            pq.offer(a + b);
        }
        return cost;
    }
}

In Java, the Queue.offer() is similar to Queue.add(). Both methods will add an element to the queue. However, when the operation fails, the offer method will return false while the add() method will raise an exception.

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

–EOF (The Ultimate Computing & Technology Blog) —