Hints: Learning to Paint

Hint 1 ↕

Come up with a DP solution that works in $O(N^3 \cdot Log(N))$. We will later optimize it to $O(N^2 \cdot Log(N))$

Hint 2 ↕

Define $indp[i]$ to be the multiset of the k best beauty of the prefix a[0...i] such that the i-th element is necessarily colored red in all those colorings.

Similarly, define $outdp[i]$ to be the multiset of the k best beauty of the prefix a[0...i] such that i-th element is necessarily colored white in all those colorings.

What are the transitions?

Hint 3 ↕

outdp[i] = Union(indp[i - 1], outdp[i - 1])

for j < i:
    indp[i] = Union(outdp[j] + a[j + 1][i])

Here, Union implies the merging of 2 sets, while ensuring that the final size of the set is at max k (by removing smaller outdated elements)

Code

Hint 4 ↕

The most expensive operation happens while updating $indp[i]$. We iterate over all the multisets of previous elements, add a delta and take the best k values out of them.

How do you efficiently extract k maximum elements out of n heaps?

Consider this famous interview problem https://leetcode.com/problems/merge-k-sorted-lists/description/

Hint 5 ↕

Store all the iterators of the indices j in a heap, along with the delta that we need to apply. Then, repeatedly pick out the maximum element after the applied delta, and put back the iterator after incrementing it.

Code