PROBLEM LINK:
Setter: Tejas Pandey
Tester: Samarth Gupta
Editorialist: Ajit Sharma Kasturi
DIFFICULTY:
MEDIUM
PREREQUISITES:
Dynamic programming, Expected value, Modular multiplicative inverse
PROBLEM:
We are given an array of length N and it is shuffled. In each turn Alice picks an index which is not chosen yet and shows it to Bob. If the number in that index is larger than all the numbers taken till now then Bob adds that number to his array. We need to find out the expected length of Bob’s array modulo 998244353.
QUICK EXPLANATION:
-
Let us first sort the array. Let dp[i] denote the expected length of Bob’s array if the first chosen index is i.
-
This can be handled in two cases. If A_i = A_i+1, then dp[i] = dp[i+1] else dp[i] = 1 + \frac{dp[i+1] +dp[i+2]+\dots +dp[N]}{N-i} .
-
Since we can pick every starting index i with probability \frac{1}{N}, the final answer will be \frac{dp[1]+dp[2]+\dots+dp[N]}{N} .
EXPLANATION:
First let us sort the array as it makes our approach simpler and doesn’t change the final solution of the problem. This problem can be solved by using dynamic programming. For this we have to define a dp state. Let dp[i] denote the expected length of Bob’s array if the first chosen index is i.
The base case is dp[N] =1 since all the remaining elements are \leq A_N and thus cannot be further added to Bob’s array.
Now let us iterate from N-1 to 1 and calculate the dp states.
Let us suppose we are at some index i. Now we need to handle the following two cases:
Case 1: A_i = A_{i+1}
- In this case, the number of elements greater than A_i is the same as that for A_{i+1} and thus the possible final Bob’s array states must also be same. Therefore, dp_i = dp_{i+1} .
Case 2: A_i \neq A_{i+1}
-
In this case suppose the first element we pick is A_i. The number of elements greater than A_i is N-i.
-
Therefore, for the next element, we could pick index {i+1} with probability \frac{1}{N-i}, index {i+2} with probability \frac{1}{N-i}, \dots, index {N} with probability \frac{1}{N-i} .
-
Thus, dp[i] = 1 + \frac{dp[i+1] +dp[i+2]+\dots +dp[N]}{N-i} .
Once we have our dp values, our final answer will be \frac{dp[1]+dp[2]+\dots+dp[N]}{N} . This is because we can pick the starting index as index 1 with probability \frac{1}{N}, index 2 with probability \frac{1}{N}, \dots, index N with probability \frac{1}{N} .
TIME COMPLEXITY:
O(N\log N) or O(N\log N + N\log MOD) per testcase depending on the implementation for calculating modular multiplicative inverse.
SOLUTION:
Editorialist's solution
#include <bits/stdc++.h>
#define int long long int
using namespace std;
int power(int x, int y, int MOD)
{
if (y == 0)
return 1;
int temp = power(x, y / 2, MOD);
temp = (temp * temp) % MOD;
if (y & 1)
temp = (temp * x) % MOD;
return temp;
}
int calcInverse(int x, int MOD)
{
return power(x, MOD - 2, MOD);
}
int32_t main()
{
int tests;
cin >> tests;
while (tests--)
{
int MOD = 998244353;
int n;
cin >> n;
vector<int> a(n + 1);
for (int i = 1; i <= n; i++)
cin >> a[i];
sort(a.begin(), a.end());
vector<int> dp(n + 1);
int sum = 0;
for (int i = n; i >= 1; i--)
{
if (i < n && a[i] == a[i + 1])
dp[i] = dp[i + 1];
else
dp[i] = (1 + sum * calcInverse(n - i, MOD)) % MOD;
sum += dp[i];
sum %= MOD;
}
int ans = (sum * calcInverse(n, MOD)) % MOD;
cout << ans << endl;
}
return 0;
}
Setter's solution
#include <bits/stdc++.h>
#define mod 998244353
#define maxn 500007
#define ll long long int
using namespace std;
ll inv[maxn];
long long dp[maxn];
int val[maxn];
ll mpow(ll a, ll b){
ll res = 1;
while(b){
if(b&1) res *= a,res %= mod;
a *= a;
a %= mod;
b >>= 1;
}
return res;
}
void pre(){
inv[0] = inv[1] = 1;
for(int i = 2;i < maxn;i++){
inv[i] = mpow(i, mod - 2);
}
}
int main() {
//freopen("inp8.in", "r", stdin);
//freopen("out8.out", "w", stdout);
pre();
int t;
cin >> t;
assert(t > 0 && t <= 100000);
while(t--) {
int n;
cin >> n;
assert(n > 0 && n <= 500000);
for(int i = 0; i < n; i++) {
cin >> val[i];
assert(val[i] > 0 && val[i] <= 500000);
}
sort(val, val + n);
reverse(val, val + n);
dp[0] = 1;
long long sm = 1;
for(int i = 1; i < n; i++) {
if(val[i] == val[i - 1])
dp[i] = dp[i - 1];
else
dp[i] = (1 + (sm*inv[i])%mod)%mod;
sm += dp[i];
sm %= mod;
}
cout << ((sm*inv[n])%mod) << "\n";
}
}
Tester's solution
#include <bits/stdc++.h>
using namespace std;
long long readInt(long long l, long long r, char endd) {
long long x=0;
int cnt=0;
int fi=-1;
bool is_neg=false;
while(true) {
char g=getchar();
if(g=='-') {
assert(fi==-1);
is_neg=true;
continue;
}
if('0'<=g&&g<='9') {
x*=10;
x+=g-'0';
if(cnt==0) {
fi=g-'0';
}
cnt++;
assert(fi!=0 || cnt==1);
assert(fi!=0 || is_neg==false);
assert(!(cnt>19 || ( cnt==19 && fi>1) ));
} else if(g==endd) {
if(is_neg) {
x=-x;
}
assert(l<=x&&x<=r);
return x;
} else {
assert(false);
}
}
}
string readString(int l, int r, char endd) {
string ret="";
int cnt=0;
while(true) {
char g=getchar();
assert(g!=-1);
if(g==endd) {
break;
}
cnt++;
ret+=g;
}
assert(l<=cnt&&cnt<=r);
return ret;
}
long long readIntSp(long long l, long long r) {
return readInt(l,r,' ');
}
long long readIntLn(long long l, long long r) {
return readInt(l,r,'\n');
}
string readStringLn(int l, int r) {
return readString(l,r,'\n');
}
string readStringSp(int l, int r) {
return readString(l,r,' ');
}
void readEOF(){
assert(getchar()==EOF);
}
#define mod 998244353
int ksm(int a, int b){
int ans = 1;
while(b){
if(b%2)
ans = ans*1ll*a%mod;
b = b/2;
a = a*1ll*a%mod;
}
return ans;
}
int inv(int a){
return ksm(a, mod - 2);
}
int main() {
// your code goes here
int t = readIntLn(1, 2e4);
int sum = 0;
while(t--){
int n = readIntLn(1, 5e5);
sum += n;
assert(sum <= 5e5);
map<int, int> mp;
for(int i = 1 ; i <= n ; i++){
int x;
if(i == n)
x = readIntLn(1, 5e5);
else
x = readIntSp(1, 5e5);
mp[x]++;
}
vector<int> vec;
for(auto j : mp)
vec.push_back(j.second);
int m = vec.size();
int dp[m];
dp[m - 1] = 1;
long long sum = dp[m - 1]*vec[m - 1];
int len = vec[m - 1];
for(int i = m - 2 ; i >= 0 ; i--){
dp[i] = (1 + sum*1ll*inv(len)%mod)%mod;
len += vec[i];
sum += dp[i]*1ll*vec[i]%mod;
sum %= mod;
}
sum = sum*1ll*inv(n)%mod;
cout << sum << '\n';
}
readEOF();
return 0;
}
Please comment below if you have any questions, alternate solutions, or suggestions.