KBOXES - Editorial

editorial
, , ,
1

PROBLEM LINK:

Practice
Contest: Division 1
Contest: Division 2
Contest: Division 3
Contest: Division 4

Author: raysh07
Tester: sushil2006
Editorialist: iceknight1093

DIFFICULTY:

Simple

PREREQUISITES:

Sorting, priority queues/sets

PROBLEM:

There are N thieves, the i-th thief has level A_i and B_i coins. All levels are distinct.
Each thief can steal from at most K other thieves; and stealing is only possible if the victim’s level is lower.
Find the maximum number of coins each thief can steal, independently.

EXPLANATION:

For a fixed thief i, what to do is fairly obvious: he can only steal from a thief with lower level, so we only look at those thieves j for whom A_j \lt A_i, and take the largest K values of B_j among them (or just take everything, if there are less than K).

What we need to figure out is how to do this quickly for all the thieves, without degenerating to quadratic complexity.
Usually, this boils down to not recomputing too much information.

Suppose we’ve found the answer for the thief with level x, i.e. we know the top K coin values among the thieves with smaller levels.
Let \{c_1, c_2, \ldots, c_K\} be the coin values with us.
What changes when we compute the answer for level x+1 instead?

Not a lot, really.
In particular, observe that:

  • When looking at the thieves with levels \lt x, the best we can do is known: it’s just \{c_1, c_2, \ldots, c_K\}.
  • That means the only thing that could possibly change is if the thief with level x has a large number of coins.
    If we want to include him, we must discard one of the existing K thieves; clearly it’s best to discard \min(c_1, c_2, \ldots, c_K).

So, when moving from level x to level x+1, the set of optimal thieves to rob from really doesn’t change much: at best, we include one new thief and discard one existing thief.

Let S be a data structure that stores the set of optimal thieves to steal from for the current level.
As noted above, it’s quite easy to keep S updated when we process thieves in increasing order of level. So, that’s exactly what we’ll do!

  • Process thieves in increasing order of level, i.e. level 1, 2, 3, \ldots
  • Suppose we’re processing level x. Let i be the index such that A_i = x-1. Since levels are distinct, this i is unique.
  • Then, do the following:
    • Insert B_i into S.
    • If S now has size larger than K, discard the smallest element in it.
  • The answer for level x is then the sum of elements in x.
    Since we’re only performing one insertion/deletion at a time, it’s easy enough to just store the sum of S in a separate variable and keep it updated with each change, to avoid having to iterate through S each time.

What data structure should we choose for S?
Looking at what we need, we see that it should be able to support quickly inserting elements, and deleting the smallest element.
This can be done by, for example, a set (std::set in C++, TreeSet in Java) or a priority queue.

TIME COMPLEXITY:

\mathcal{O}(N \log N) per testcase.

CODE:

Editorialist's code (PyPy3) 
for _ in range(int(input())):
    n, k = map(int, input().split())
    a = list(map(int, input().split()))
    b = list(map(int, input().split()))
    
    order = list(range(n))
    order.sort(key = lambda x: a[x])
    
    from heapq import heappush, heappop
    pq, ans, sm = [], [0]*n, 0
    for i in order:
        ans[i] = sm
        
        heappush(pq, b[i])
        sm += b[i]
        if len(pq) > k:
            sm -= pq[0]
            heappop(pq)
    print(*ans)
1 Like