Video Solution

By Abhiraj Mallangi

Note: The video solution might not be the same as other solutions. Code in Java.

Solution

Explanation

In this problem, we're asked to find the minimum cost to divide a stick with length $x$ into $n$ sticks with given lengths. It helps to work backwards; what if we start with sticks of lengths $d_1,\ldots,d_n$ and merge them into one?

It turns out that this can be solved using Huffman Coding (also see Wikipedia). The algorithm is simple; take the two shortest sticks, merge them into one, and repeat.

If you're wondering why Huffman Coding always produces an optimal solution, see here.

As we want to select both shortest sticks and then insert the combined new stick, we will want to have a data structure where we are able to get the smallest value in the collection, remove it and add new values to the collection. The insert and remove operations can be done easily with a priority queue in $\mathcal{O}(\log n)$ time, while retrieving the smallest value takes constant time.

Implementation

Time Complexity: $\mathcal{O}(N \log N)$

Copy
import heapq

x, n = map(int, input().split())
d = list(map(int, input().split()))

heapq.heapify(d)
res = 0
for _ in range(n - 1):
	a = heapq.heappop(d)
	b = heapq.heappop(d)
	res += a + b
	heapq.heappush(d, a + b)

print(res)