MEXQUER - Editorial

PROBLEM LINK:

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

Author: progokcoe
Tester: raysh07
Editorialist: iceknight1093

DIFFICULTY:

TBD

PREREQUISITES:

None

PROBLEM:

You’re given a multiset of integers, which has A_i occurrences of i.
In one move, you can choose an integer x from the multiset and replace it with either 2x or 2x+1.

For each x from 0 to N, find the minimum number of moves you need to make the multiset have a mex of x.

EXPLANATION:

For the mex of a multiset to be x, it should contain at least one occurrence of every integer from 0 to x-1 and no occurrences of x.
Elements greater than x don’t matter.

Let \text{ans}[y] denote the minimum number of moves needed to get a mex of y (or -1 if we can’t).
Note that our operation of x\to 2x or x\to 2x+1 can only increase integers.
So, to get a mex of y:

  • We need to get rid of every occurrence of y.
    Since y\to2y + 1 strictly increases the value of one occurrence, we need one such move for each of the A_y occurrences of y.
  • That leaves us to then obtain at least one occurrence of everything between 0 and y-1.

Let’s try to compute \text{ans}[y] in increasing order of y.
Clearly, \text{ans}[0] = A_0 as noted above.
For y\gt 0 however, observe that by definition, \text{ans}[y-1] - A_{y-1} is exactly the minimum number of moves needed to obtain at least one occurrence of every integer in [0, y-2].
The only thing we need to do now, is figure out how to get one occurrence of y-1.

If A_{y-1} \gt 0 already then we’re done.
Otherwise, our only hope is to choose some smaller x \lt y-1, and operate on it till it reaches y-1.
It turns out, there aren’t too many x that can reach y-1 at all - in fact, there are only \mathcal{O}(\log y) of them, so we can simply try all of them!

How?

For any integer y, the only integer that can reach y in a single move is \left\lfloor \frac{y}{2} \right\rfloor.
But then, by the same logic, the only integer that can reach \left\lfloor \frac{y}{2} \right\rfloor in a single move is \left\lfloor \frac{\left\lfloor \frac{y}{2} \right\rfloor}{2}\right\rfloor = \left\lfloor \frac{y}{4} \right\rfloor, and so on.

More concretely, the only integers that can possibly reach y are those of the form \left\lfloor \frac{y}{2^k} \right\rfloor for some k.
However, we only need to consider k \leq \log_2 y since for anything larger, the floor of the division is 0 anyway.
This means we need to check for at most \log_2 y values of x.


So, check for each x\leq y-1 that can be turned into y whether there is some ‘extra’ occurrence of it (i.e, whether A_x \gt 1), and if so, convert one such occurrence to y-1.
If there are multiple valid x, it’s optimal to choose the largest one, since it’ll need the minimum number of steps.
(Note that it’s technically possible to choose some x with A_x = 1 and convert it to y-1, then attempt to recreate x using a smaller number. However, this will need the same number of moves as just taking that smaller number to y-1 directly, so we don’t need to explicitly consider such a case.)

This way, we can obtain \text{ans}[y] from \text{ans}[y-1] in \mathcal{O}(\log y) time, so all answers can be computed in \mathcal{O}(N\log N) time in total.

TIME COMPLEXITY:

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

CODE:

Author's code (C++)
#include <iostream>
#include <vector>
#define ll  long long int

using namespace std;

int main()
{
	ll t,n;cin>>t;
	while(t--)
	{
		cin>>n;
		vector <ll> COUNT(n+1,0),freq(n+1,0);
		for(ll i=0;i<n;i++)
		{
			cin>>COUNT[i];
			freq[i] = COUNT[i];
		}
		vector <ll> dp(n,0);
		// dp[i] = minimum number of steps required to make frequency of all elements 0 to i greater than 0
		// dp[i] = -1  if not possible to achieve above mentioned

		// base case
		if(freq[0]>0) dp[0]=0;
		else          dp[0]=-1;

		for(ll i=1;i<n;i++)
		{
			if(dp[i-1]==-1) dp[i]=-1;
			else
			{
               if(freq[i]>0) dp[i]=dp[i-1];
               else 
               {
               	   dp[i]=dp[i-1];
                   ll x=i;
                   while(freq[x]==0)
                   {
                   	  if(x == 0) {dp[i]=-1;break;}
                   	  freq[x]++;
                   	  dp[i]++;
                   	  x/=2;
                   	  freq[x]--;
                   }
               }
			}
		}
        for(ll i=0;i<n;i++)
        {
        	if(i == 0) cout<<COUNT[0]<<' ';
        	else       
        	{
        		if(dp[i-1]==-1) cout<<-1<<' ';
        		else            cout <<dp[i-1] + COUNT[i]<<' ';
        	}
        }
        cout<<dp[n-1]<<"\n";
	}
}
Tester's code (C++)
#include <bits/stdc++.h>
using namespace std;
#define int long long
#define INF (int)1e18
#define f first
#define s second

mt19937_64 RNG(chrono::steady_clock::now().time_since_epoch().count());

void Solve() 
{
    int n; cin >> n;
    
    vector <int> f(n + 1, 0);
    for (int i = 1; i <= n; i++) cin >> f[i - 1];
    
    int mex = 0;
    while (f[mex]) mex++;
    
    for (int i = 0; i < mex; i++){
        cout << f[i] << " ";
    }
    
    cout << 0 << " ";
    
    int ans = 0;
    
    for (int i = mex; i < n; i++){
        if (!f[i]){
        int x = i / 2;
        int c = 1;
        while (true){
            if (f[x] > 1){
                f[x]--;
                ans += c;
                break;
            }
            c++;
            if (x == 0){
                ans = -INF;
                break;
            }
            x /= 2;
        }
        }
        
        if (ans < 0) cout << -1 << " ";
        else cout << ans + f[i + 1] << " ";
    }
    cout << "\n";
}

int32_t main() 
{
    auto begin = std::chrono::high_resolution_clock::now();
    ios_base::sync_with_stdio(0);
    cin.tie(0);
    int t = 1;
    // freopen("in",  "r", stdin);
    // freopen("out", "w", stdout);
    
    cin >> t;
    for(int i = 1; i <= t; i++) 
    {
        //cout << "Case #" << i << ": ";
        Solve();
    }
    auto end = std::chrono::high_resolution_clock::now();
    auto elapsed = std::chrono::duration_cast<std::chrono::nanoseconds>(end - begin);
    cerr << "Time measured: " << elapsed.count() * 1e-9 << " seconds.\n"; 
    return 0;
}
Editorialist's code (Python)
for _ in range(int(input())):
    n = int(input())
    a = list(map(int, input().split()))
    a.append(0)
    ans = [0]*(n+1)
    ans[0] = a[0]
    for i in range(1, n+1):
        if ans[i-1] == -1:
            ans[i] = -1
            continue
        
        ans[i] = ans[i-1] - a[i-1] + a[i]
        if a[i-1] == 0:
            x = i-1
            steps = 0
            while x > 0:
                x //= 2
                steps += 1
                if a[x] > 1:
                    a[x] -= 1
                    a[i-1] += 1
                    ans[i] += steps
                    break
        if a[i-1] == 0: ans[i] = -1
    print(*ans)
1 Like