Maximum Prefix Mex Sum - Editorial

About the Problem

Setter(s) Divyansh Verma
Tester(s) Vichitr Gandas, Mayank Pugalia
Difficulty Easy
Topics Observation, Greedy, Maths, Combinatorics
Practice Link
Contest Link

Solution Idea

The prefix Mex sum is maximum when we put all unique numbers in increasing order at the start. Once a number is missed from 0,1,2..., this chain breaks and then Mex will remain same for rest of the prefix arrays. Example:
A=[0,1,2,5,7]
Here Mex for first three prefix arrays is 1,2,3 respectively. After that we see that 3 is not present. Hence for rest of subarrays Mex will remain 3 only.
This way, the sum is maximized when we arrange initial unique numbers in order 0,1,2... like this. And if this order breaks or we are left with duplicate items, we can permute them in any way. Here in the beginning, each unique number x can be picked with count[x] ways where count[x] is number of times it occurs in the array. And let’s say number of remaining elements after break is y, they can be arranged in y! ways. Hence total number of permutations with max prefix Mex sum would be count[0]*count[1]*...*count[k]*y! where k is first number where this sequence 0,1,2,... breaks.

Count can be maintained with map. We can reduce the time complexity here by using a vector instead because the sequence if continued won’t exceed n and frequency of rest of the items is not needed, because they will be part of left over items which can be permuted.

Complexity Analysis

Time Complexity: \mathcal{O}(n \log{n}) for pre calculation and to find the answer which can be reduced to \mathcal{O}(n) using a vector instead of map.
Space Complexity: \mathcal{O}(n) for maintaining the frequency.

Codes

Setter's Code
#include<bits/stdc++.h>
using namespace std;
#define int long long int

const int MAXN = 1e5 + 10;
const int MOD = 1e9 + 7;

int fact[MAXN];

void preCompute(){
    fact[0] = 1;
    for(int i=1;i<MAXN;i++) fact[i] = (fact[i-1]*i)%MOD;
}


class FindMexMax {
public:
    int getOutput(vector<int> a){

        map<int,int> which;
        for(int v:a) which[v]++;
        
        int cnt = a.size(),ans = 1;

        for(int i=0;i<a.size();i++){
            if(which[i] == 0) break;
            ans = (ans*which[i])%MOD;
            which[i]--;
            cnt--;
        }

        ans = ans*fact[cnt]%MOD;
        return ans;
    }
};

signed main(){
  	preCompute();
  int t;
  cin>>t;
  while(t--){
	int n; cin>>n;
  	vector<int> a(n);
    for(int &v:a) cin>>v;
    FindMexMax fMM;
    cout<<fMM.getOutput(a)<<endl;
  }
}
Tester's Code
/***************************************************

@author: vichitr
Compiled On: 13th Mar 2021

*****************************************************/
#include<bits/stdc++.h>
#define endl "\n"
#define int long long
using namespace std;

const int MOD = 1e9 + 7;
const int N = 1e5 + 7;
int fact[N];

void init(){
	fact[0] = 1;
  	for(int i=1;i<N;i++){
		fact[i] = (fact[i-1] * i) % MOD;
	}
}

void solve(){
	int n; cin>>n;
  	assert(n <= 1e5);
  	int a[n];
  	map<int, int> M;
  	for(int i=0;i<n;i++){
		cin>>a[i];
	  	assert(a[i] >= 0 and a[i] <= 1e9);
	  	M[a[i]]++;
	}
  	int ans = 1;
  	int tot = 0;
  	for(int i = 0; i < n; i++){
		if(M[i] == 0){
			break;
		}
	  	ans *= M[i];
	  	ans %= MOD;
	  	tot++;
	}
  	ans *= fact[n - tot];
  	ans %= MOD;
  	cout<<ans<<'\n';
}

signed main(){
	int t = 1;
  	cin>>t;
  	assert(t <= 5);
  	init();
  	while(t--){
		solve();
	}
  	return 0;
}

If you have used any other approach, share your approach in comments!
If you have any doubts, ask them in comments!

1 Like