Collect Every Pokémon
National Olympiad — National round (2024 season)
Submission name: pokemons.cpp
Problem summary
Each of m days independently chooses one of n Pokémon types. Count the length-m sequences that contain every type at least once, modulo 1,102,024,631.
Input
m n.
Output
The number of sequences modulo 1,102,024,631.
Constraints
1 ≤ m ≤ 10¹⁸; 1 ≤ n ≤ 10⁷.
Worked example
An independently prepared example for this explanation.
Input
3 2
Output
6
Solution analysis
- By inclusion–exclusion, the answer is the sum over i=0…n of (−1)^(n−i)·C(n,i)·i^m. Here i is the number of types allowed to appear. If m<n the answer is immediately zero. The prime modulus exceeds n, so every denominator in the binomial recurrence is invertible.
- Computing every i^m independently would repeat many exponentiations. The function f(i)=i^m modulo the prime is completely multiplicative: f(ab)=f(a)f(b). A linear prime sieve computes f(p) by binary exponentiation only when p is prime, then fills each composite value once by multiplying its two already known factors. This saves most exponentiation work.
- Compute modular inverses in linear time using inv[i]=mod−(mod/i)·inv[mod%i] modulo mod. Starting with C(n,0)=1, obtain consecutive binomial coefficients by multiplying (n−i)/(i+1). For m=3,n=2 the formula is 2³−2·1³=6. The i=0 contribution vanishes because m is positive.
Why it works
Inclusion–exclusion cancels every sequence missing any type and retains every sequence using all types exactly once. Complete multiplicativity follows from (ab)^m=a^mb^m; the linear sieve assigns each composite using its least-prime-factor representation, so its stored powers are correct. The inverse and binomial recurrences give the exact modular coefficients because the modulus is prime and all denominators are smaller. Thus the final sum equals the desired count.
Watch out. The modulus is 1,102,024,631, not 10⁹+7. Reduce 64-bit products before storing them. Do not exponentiate every integer separately at the full bound.
Commented C++17 solution
An independently written complete program with standard headers.
#include<iostream>
#include<vector>
// Algorithm: By inclusion–exclusion, the answer is the sum over i=0…n of (−1)^(n−i)·C(n,i)·i^m.
using namespace std;
const long long mod=1102024631;
int power(long long base,long long exponent){
long long result=1;
while(exponent>0){
if(exponent&1){
result=result*base%mod;
}
base=base*base%mod;
exponent>>=1;
}
return static_cast<int>(result);
}
int main(){
ios::sync_with_stdio(false);
cin.tie(nullptr);
long long days;
int n;
cin>>days>>n;
if(days<n){
cout<<"0\n";
return 0;
}
vector<int>least(n+1),powers(n+1),primes;
powers[1]=1;
for(int i=2;i<=n;++i){
if(least[i]==0){
least[i]=i;
primes.push_back(i);
powers[i]=power(i,days);
}
for(int prime:primes){
if(prime>least[i]||prime>n/i){
break;
}
least[i*prime]=prime;
powers[i*prime]=static_cast<int>(1LL*powers[i]*powers[prime]%mod);
}
}
vector<int>inverse(n+1);
inverse[1]=1;
for(int i=2;i<=n;++i){
inverse[i]=static_cast<int>(mod-(mod/i)*inverse[mod%i]%mod);
}
long long choose=1;
long long answer=0;
for(int i=0;i<=n;++i){
long long term=choose*powers[i]%mod;
if((n-i)&1){
answer=(answer-term+mod)%mod;
}
else{
answer=(answer+term)%mod;
}
if(i<n){
choose=choose*(n-i)%mod*inverse[i+1]%mod;
}
}
cout<<answer<<'\n';
}