RTAR - Editorial

PROBLEM LINK:

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

Author: pols_agyi_pols
Tester: mexomerf
Editorialist: iceknight1093

DIFFICULTY:

TBD

PREREQUISITES:

Sorting, sets/priority queues

PROBLEM:

There are N programmers, the i-th has A_i stars and is ratist iff S_i = 1.
A ratist programmer will refuse to sit next to anyone with strictly less stars.
Find the maximum possible number of programmers who can be seated around a round table.

EXPLANATION:

We’re clearly only constrained by the ratist programmers, so we need to focus on them.

Let r_1 \gt r_2 \gt \ldots \gt r_k be distinct ratings of ratist programmers. We’ll first try to characterize when they can be arranged around a table (and the maximum number of non-ratists that can be chosen).
(It’s enough to consider distinct ratings because if we’re able to find an arrangement with just these, others with the same rating can be placed next to them without affecting the arrangement.)

For now, we’ll deal with k \geq 2; the smaller values will be considered at the end.

Clearly, no two ratist programmers can be seated next to each other - the one with more stars will become unhappy.
So, we need at least k non-ratist ‘separators’.

The separators must have large enough values to be valid. Specifically,

  • We need at least two non-ratists with values \geq r_1, to place on its left and right.
  • One side of r_2 can be covered with one of the above, the other side requires another separator with value \geq r_2.
  • Similarly, we need another separator with value \geq r_3, and so on till we reach r_{k-1}.
  • r_k doesn’t need anything special, since it’ll be covered by the separators on its sides anyway.

More generally, we require one non-ratist with value \geq r_i for every 1 \leq i \lt k, and then one extra separator with value \geq r_1.
Further, observe that as long as we have at least one more non-ratist with a value \geq r_k, we can in fact choose every non-ratist: place the extra non-ratist next to r_k, after which we’ll have two non-ratists adjacent to each other and all the other non-ratists can be placed between them in any order.
On the other hand, if we don’t have this extra, we cannot place any more non-ratists at all.

The remaining cases are:

  • If k = 0, meaning there are no ratists chosen, we can simply take all non-ratists.
  • If k = 1, there are a couple of possibilities:
    • If there are at least two non-ratists with values \geq r_1, we can take all non-ratists.
    • Otherwise, we can only take the non-ratists with values \geq r_1 (of which there are either 0 or 1).

Now, let’s move on to solving the problem.

We’ll iterate over all people in descending order of their A_i values, and process all people of a certain value at the same time.
Along the way, we’ll also maintain the current best answer as a sort of open-ended ‘arc’; which we’ll try to close into a circle as we go. The arc will be bordered on both sides with non-ratists.

Suppose we’re processing people with x stars.
Let there be b_x ratists among them.

Then,

  • If b_x \gt 0, we can try to close the arc using ratists of rating b_x to obtain a circle.
    Note that this is possible only if the arc contains at least two non-ratists already.
    • If the arc is closed with these b_x people to form a circle, there are two possibilities: either no more people can be added to it, or all non-ratists can be added to it.
      As discussed earlier, this depends on whether we have any extra non-ratists or not.
      Compute the size of the completed circle appropriately and maximize the answer with it.
    • We have another option available to us: close the circle with these b_x people, then remove some ratists from the circle - this will allow for all non-ratists to be chosen.
      Here, it’s clearly best to remove whichever value has the least frequency from the circle.
  • Then, we attempt to update the arc with people having x stars.
    • If the arc doesn’t exist yet, you can create it using these b_x people and any available non-ratists; but only if at least two non-ratists are available.
    • If there’s an extra non-ratist available, we can extend the arc with all b_x ratists.
    • Otherwise, we can possibly replace some existing block of ratists with this block of b_x people.
      It’s obviously best to choose the smallest existing block to replace in this case.

To do this quickly, note that we only need a few pieces of knowledge: the current size of the arc, the minimum size of a block of ratists within the arc, and the number of non-ratists available to us.
The counts are simple enough to maintain, and the sizes of the people in the arc are easily maintained and updated using a data structure like a (multi)set or priority queue.

This gives us a solution in \mathcal{O}(N\log N) overall.

Don’t forget the case where no ratists are chosen, i.e, just the count of zeros in S.

TIME COMPLEXITY:

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

CODE:

Author's code (C++)
#include <bits/stdc++.h>
using namespace std;
#define ll long long

int main(){
	ios_base::sync_with_stdio(false);
	cin.tie(NULL);
    ll kitne_cases_hain;
    kitne_cases_hain=1;
    cin>>kitne_cases_hain;
    while(kitne_cases_hain--){          
        ll n;
        cin>>n;
        ll a[n];
        ll ans=0;
        for(int i=0;i<n;i++){
            cin>>a[i];
        }
        string s;
        cin>>s;
        map <pair<ll,ll>,ll> m; 
        for(int i=0;i<n;i++){
            m[{-1*a[i],s[i]-'0'}]++;
        }
        ll cnt=0;
        ll total=0;
        multiset <ll> ms;
        ll x;
        for(auto it:m){
            ans=max(it.second,ans);
            if(it.first.second==0){
                total+=it.second;
                cnt+=it.second;
                ans=max(ans,total);
            }else{  
                if(cnt!=0){
                    cnt--;
                    total+=it.second;
                    ans=max(ans,total);
                    ms.insert(it.second);
                    if(cnt==0){
                        cnt++;
                        x=*ms.begin();total-=x;
                        ms.erase(ms.begin());
                    }
                }
            }
        }
        cout<<ans<<"\n";
    }
	return 0;
}
Editorialist's code (Python)
import heapq
for _ in range(int(input())):
    n = int(input())
    a = list(map(int, input().split()))
    s = input()
    indices = list(range(n))
    indices.sort(key = lambda x: a[x])
    
    zeros = s.count('0')
    ans = good = badsm = store = 0
    p = n-1
    pq = []
    while p >= 0:
        i, q = indices[p], p
        curbad = 0
        while q >= 0:
            j = indices[q]
            if a[i] != a[j]: break
            curbad += s[j] == '1'
            store += s[j] == '0'
            q -= 1
        
        if curbad > 0:
            if good <= 1:
                if good + store >= 2: ans = max(ans, curbad + zeros)
                else: ans = max(ans, curbad + good + store)
            else:
                if len(pq) + 1 == good and store == 0:
                    ans = max(ans, badsm + curbad + good)
                    ans = max(ans, badsm + curbad + zeros - pq[0])
                else: ans = max(ans, badsm + curbad + zeros)

            good += store
            store = 0
            if good >= 2:
                badsm += curbad
                heapq.heappush(pq, curbad)
                while len(pq) > good - 1:
                    badsm -= heapq.heappop(pq)
        
        p = q
    print(max(ans, zeros))