Problem statement

https://leetcode.com/problems/minimum-cost-to-connect-sticks/

Solution

Actually, what is asked to do in this problem is to construct classical Huffman Tree, where we have the following minimization problem.

\(\min\limits_{trees} \sum\limits_x f(x) \cdot d(x)\) where $d(x)$ are depths of leaf in our tree and $f(x)$ are weights of sticks. We can note that each stick will be paid for exaclty $d(x)$ times. It is known that the optimal tree can be constructed with greedy strategy: we will take always two smallest weights and merge them to one node.

Complexity

Time complexity is O(n log n), space complexity is O(n).

Code

class Solution:
    def connectSticks(self, sticks):
        heapify(sticks)
        ans = 0
        for _ in range(len(sticks) - 1):
            a, b = heappop(sticks), heappop(sticks)
            ans += (a + b)
            heappush(sticks, a + b)

        return ans