KMO

Tree interpretation

Lengths $b_i$ of a prefix code over an alphabet of size $K$ satisfy Kraft's inequality $\sum K^{-b_i}\leq1$, and conversely every such integer length sequence can be realized. To see the converse, process lengths in increasing order, expanding each unused position into $K$ positions at the next depth and placing the required leaves.

Thus we seek a rooted tree with at most $K$ children per internal node, minimizing the sum of leaf weight times leaf depth. Zero-weight input elements remain real leaves and must not be discarded.

Huffman merging

If $N=1$, output $0$. Otherwise a full $K$-ary tree has a leaf count congruent to $1$ modulo $K-1$. Add

$$d=(K-1-(N-1)\bmod(K-1))\bmod(K-1)$$

dummy leaves of weight zero. This preserves the optimum: remove unary nodes from an optimal tree, and consolidate incomplete internal nodes at the deepest level until only one incomplete node remains. Moving leaves toward the root cannot increase the cost for nonnegative weights. Filling its missing positions requires exactly these $d$ dummy leaves.

There is an optimal full tree whose $K$ lightest leaves are deepest siblings. Exchanging a lighter shallow leaf with a heavier deep leaf cannot increase cost. Contract these siblings into one leaf whose weight is their sum. The original cost equals the contracted cost plus that sum. Repeating this operation gives $K$-ary Huffman coding.

Avoid constructing dummy leaves explicitly. Merge

$$r=K-d=2+(N-2)\bmod(K-1)$$

real nodes first, and merge exactly $K$ nodes in every later step. In particular, $K>N$ gives $r=N$. If $N-1$ is divisible by , the first merge must also use nodes.

Use a min-heap. Record the newly created node as the parent of each removed node. Parent indices are larger than child indices, so compute depths in reverse node order without recursion. Output the original leaves' depths in input order.

Bounds and complexity

There are at most $N-1$ internal nodes, each with at least two children. Thus every leaf depth is at most $N-1$. Time is $O(N\log N)$ and space is $O(N)$.

A merged weight is at most $\sum a_i\leq2\cdot10^{14}$. Assigning every leaf depth $\lceil\log_K N\rceil$ is feasible, so the optimum is at most . Signed 64-bit integers suffice; never compute directly.

Subtasks

For subtask 1, sort weights in decreasing order and enumerate nondecreasing depths from $0$ through $N-1$. An exchange argument guarantees an optimum with this ordering. There are $\binom{2N-1}{N}$ sequences, each checked in , giving time.

Subtask 2 is binary Huffman coding: repeatedly merge the two smallest weights in $O(N\log N)$ time.

Exact certificate checking

Count leaves at each depth. Start with one position at depth zero, subtract leaves placed there, and multiply remaining positions by $K$ before processing the next depth. Reject if a depth needs too many positions. Cap the position count at $N+1$: there are only $N$ leaves, so larger exact counts cannot affect feasibility.

The checker recomputes the optimum and uses it as a remaining cost budget. Before subtracting $a_i b_i$, use division to detect an exceeded budget. This avoids overflow even for incorrect contestant output.

Solution written by GPT6