E - Kth Takoyaki Set Editorial by spheniscine


We need a data structure that models a set \(S\) and supports the following operations:

  • find and remove the minimum element of \(S\)
  • add an element to \(S\) (unless it already is in \(S\))

Most languages have some data structure that supports these operations in \(O(\log n)\) where \(n\) is the number of elements currently in the set, using a balancing binary search tree (BBST). Examples include std::set in C++, std::collections::BTreeSet in Rust, and java.util.TreeSet for Java and Kotlin.

We initialize \(S\) with a single element, \(0\). We can then solve this problem by doing the following \(K\) times:

  • find and remove the minimum element of \(S\), let this be \(x\)
  • for each \(A_i\), add \(x + A_i\) to \(S\).

The answer is then the minimum element of \(S\).

You could prove that this is correct by the following outline: let \(T\) be the infinite set of possible prices (considering \(0\) a price as well), and note that our algorithm, assuming unlimited time and space (and no integer overflow), will eventually find any element of \(T\). Note that after you withdraw an element from \(S\), you will never reinsert another element less than or equal to it. Thus, the first \(i\) elements withdrawn from \(S\) is equal to the smallest \(i\) elements in \(T\). The final answer is thus the \(K+1\)-st element of \(T\), which is what we want since the original problem does not count \(0\) as a price.

Thus the problem is solved in time \(O (NK \log NK)\) and space \(O(NK)\), which should pass.

Note that the intermediate and final values may not fit in a standard 32-bit signed integer data type; a 64-bit signed integer should be sufficient.

posted:
last update: