PROBLEM LINK:
Practice
Div-2 Contest
Div-1 Contest
Author: Ildar Gainullin
Tester: Alexander Morozov
Editorialist: Mohammed Ehab
DIFFICULTY:
MEDIUM-HARD
PREREQUISITES:
Dynamic Programming
PROBLEM:
Given an array A of length N, let’s take constants X_1, X_2, \ldots, X_N such that X_1 \cdot A_1+X_2 \cdot A_2+ \ldots X_N \cdot A_N is a multiple of 239. You want to count the number of ways to pick these constants such that each of them is between 0 and 238 (inclusive,) and the constants are distinct.
QUICK EXPLANATION:
Let dp_s denote the answer if all the elements outside of the subset s are replaced by one element equal to their sum, let’ call it “the combined element.” To calculate the dp, you can count the number of ways to make the coefficients of the elements in s distinct without caring about the coefficient of the combined element. Then, you can subtract the number of ways where the combined element’s coefficient is equal to one of the other coefficients. That gives you 2 cases: if the combined element is non-zero, you can pick the rest of the coefficients as you want, and then use it to make the expression divisible by $239, so dp_s=\frac{239!}{(239-|s|)!}-\sum_{j \in s} dp_{s \setminus j}. If it’s zero, then you can remove it from the array, solve without it, and then multiply the answer by 239-|s|. Let j be any index that belongs to s. Solving without the combined element is equivalent to adding j to it (since its sum is 0,) so dp_s=(239-|s|) \cdot dp_{s \setminus j}.
EXPLANATION:
Let’s forget about the constraint that the constants should be distinct for a second. Let’s pick the first N-1 constants arbitrarily; really, choose anything. I claim that if A_N \neq 0, there’s one and only one way to pick X_N so that the sum is a multiple of 239. Why is that? Notice that we want:
X_1 \cdot A_1+X_2 \cdot A_2+ \ldots X_N \cdot A_N \equiv 0 \pmod{239}
X_N \cdot A_N \equiv -(X_1 \cdot A_1+X_2 \cdot A_2+ \ldots + X_{N-1} \cdot A_{N-1}) \pmod{239}
Now, since 239 is a prime, every number other than 0 has a modular inverse. If you multiply the above equation by the modular inverse of A_N, you get the one and only one possible value for X_N. So solving the problem without the distinctivity constraint (and with at least one non-zero element in A) is easy. Just pick N-1 of the constants as you please, and use the last constant to “fix” the value and make it divisible by 239.
How can we add the distinctivity condition? Maybe a first step is: instead of picking the first N-1 constants as we please, we can pick them to be distinct instead. The number of ways to do that is easy. The first constant has 239 possible values, leaving 238 values for the second constant, leaving 237 for the third, and so on. We just need to multiply these together. So that count is just \frac{239!}{(239-N+1)!}. The only problem is: the last constant (X_N) that we use to fix the expression may clash with one of the other constants. That is, maybe X_N will end up being equal to another constant X_j and ruin the solution. A good idea would be to try and count the number of ways this could happen and subtract it. That is, we wanna count the number of ways:
X_1 \cdot A_1+X_2 \cdot A_2+ \ldots X_N \cdot A_N \equiv 0 \pmod{239}
with the additional constraint that X_j=X_N for some j. Sounds harder 
However, let’s fix j. The expression turns into:
X_1 \cdot A_1+X_2 \cdot A_2+ \ldots + X_{j-1} \cdot A_{j-1} + X_{j+1} \cdot A_{j+1} + \ldots + X_N \cdot (A_j + A_N) \equiv 0 \pmod{239}
That is, we factored out X_N as a common factor between A_j and A_N. Why is that useful? This is actually the exact same problem but on a slightly different array where A_j and A_N were removed and their sum was added instead! I’ll recap how we got here. We counted the number of ways to give distinct values to the first N-1 constants. Then, we wanted to subtract the number of ways X_N will be equal to one of them. The fact that 2 of the constants are equal enabled us to merge them into one, which gave us the exact same problem on a smaller array.
If you keep repeating the same process; that is, you solve the problem recursively. You’ll find that in the expression, your elements will be split into 2 subsets: the free subset whose constants you can choose freely… and the combined subset whose elements are all multiplied by the same constant. That happens because we keep removing the last element and some other element A_j, and we insert their sum instead, provided they get multiplied by the same constant. The combined subset contains the elements we chose down the way, and the free subset contains the rest of the array.
The fact that we split the array into 2 subsets (and that N \le 20 :p) hints us to use subset dp. Let dp_s denote the number of solutions to our problem, if s is the free subset. That is, if the rest of the elements were all combined into one element equal to their sum (mod 239). What is the recurrence?
Case #1: the combined element is not 0.
We agreed we can always use it to make the expression divisible by 239, so we can start out with \frac{239!}{(239-|s|)!}, the number of ways to pick the free variables’ constants, and then we wanna subtract the number of ways the combined variables’ constant clashes with one of the free variables’ constants. Let it clash with X_j's constant. We want to subtract dp_{s \space without \space j} from dp_s. That is, we add j to the combined variables and subtract the answer. So overall:
dp_s=\frac{239!}{(239-|s|)!}-\sum_{j \in s} dp_{s \setminus j}
Case #2: the combined element is 0.
That may seem to ruin our whole idea but actually not really. Let’s remove the combined element from the array and count the number of solutions. Since the combined element is 0, you can multiply it by any constant without changing the sum. But this constant has to be different from the rest, so the number of ways to choose it is 239-|s|. What about the number of solutions if we remove the combined element? How can we express that in terms of our dp?
Get any element j \in s. Since the combined element is 0, removing it and j and inserting their sum instead is equivalent to, well, removing it. This is represented by dp_{s \space without \space j}!
So in that case, dp_s=(239-|s|) \cdot dp_{s \setminus j}
Of course, you should implement the subset dp using bitmasks.
Time complexity: O(2^N \cdot N)
SOLUTIONS:
Setter's Solution
#include <bits/stdc++.h>
using namespace std;
#pragma GCC target ("avx2")
#pragma GCC optimization ("O3")
#pragma GCC optimization ("unroll-loops")
using ll = long long;
using ld = long double;
#define mp make_pair
const int p = 998244353;
int mul(int a, int b) {
return (1LL * a * b) % p;
}
int add(int a, int b) {
int s = (a+b);
if (s>=p) s-=p;
return s;
}
int sub(int a, int b) {
int s = (a+p-b);
if (s>=p) s-=p;
return s;
}
int po(int a, int deg)
{
if (deg==0) return 1;
if (deg%2==1) return mul(a, po(a, deg-1));
int t = po(a, deg/2);
return mul(t, t);
}
int inv(int n)
{
return po(n, p-2);
}
mt19937 rnd(time(0));
/*
const int N = 1000005;
vector<int> facs(N), invfacs(N);
void init()
{
facs[0] = 1;
for (int i = 1; i<N; i++) facs[i] = mul(facs[i-1], i);
invfacs[N-1] = inv(facs[N-1]);
for (int i = N-2; i>=0; i--) invfacs[i] = mul(invfacs[i+1], i+1);
}
int C(int n, int k)
{
if (n<k) return 0;
if (n<0 || k<0) return 0;
return mul(facs[n], mul(invfacs[k], invfacs[n-k]));
}*/
/*struct DSU
{
vector<int> sz;
vector<int> parent;
void make_set(int v) {
parent[v] = v;
sz[v] = 1;
}
int find_set(int v) {
if (v == parent[v])
return v;
return parent[v] = find_set(parent[v]);
}
void union_sets(int a, int b) {
a = find_set(a);
b = find_set(b);
if (a != b) {
if (sz[a] < sz[b])
swap(a, b);
parent[b] = a;
sz[a] += sz[b];
}
}
DSU (int n)
{
parent.resize(n);
sz.resize(n);
for (int i = 0; i<n; i++) make_set(i);
}
};*/
vector<int> C(20);
const int P = 239;
int main()
{
ios_base::sync_with_stdio(0);
cin.tie(nullptr);
int n;
cin>>n;
C[0] = 1;
for (int i = 1; i<20; i++)
{
C[i] = mul(C[i-1], P + 1 - i);
}
vector<int> a(n);
for (int i = 0; i<n; i++) cin>>a[i];
vector<int> dp(1<<n);
int sum = 0;
for (auto it: a) sum+=it;
if (sum%P) dp[(1<<n)-1] = 1; else dp[(1<<n)-1] = P;
for (int mask = (1<<n)-2; mask>0; mask--)
{
int bits = 0;
sum = 0; for (int i = 0; i<n; i++) if (mask&(1<<i)) {sum+=a[i]; bits++;}
if (sum%P)
{
dp[mask] = C[n-bits];
for (int i = 0; i<n; i++) if ((mask&(1<<i))==0)
{
dp[mask] = sub(dp[mask], dp[mask^(1<<i)]);
}
}
else
{
int bit = 0;
while (mask&(1<<bit)) bit++;
dp[mask] = mul(dp[mask^(1<<bit)], P - (n-bits));
}
//cout<<mask<<' '<<dp[mask]<<endl;
}
cout<<dp[1];
}
Editorialist's Solution
#include <bits/stdc++.h>
using namespace std;
#define mod 998244353
#define MX 239
int n,a[25];
long long fact[MX+5],inv[MX+5],dp[(1<<21)];
long long pow_log(long long x,int y)
{
if (!y)
return 1;
long long ret=pow_log(x,y/2);
ret=(ret*ret)%mod;
if (y%2)
ret=(ret*x)%mod;
return ret;
}
int sum(int mask) //the sum of the elements in this subset
{
int ret=0;
for (int i=0;i<n;i++)
{
if (mask&(1<<i))
ret=(ret+a[i])%MX;
}
return ret;
}
int main()
{
scanf("%d",&n);
for (int i=0;i<n;i++)
scanf("%d",&a[i]);
fact[0]=1;
for (int i=1;i<=MX;i++)
fact[i]=(i*fact[i-1])%mod; //factorial
inv[MX]=pow_log(fact[MX],mod-2);
for(int i=MX-1;i>=0;i--)
inv[i]=((i+1)*inv[i+1])%mod; //factorial inverse
dp[0]=(sum((1<<n)-1)? 1:MX);
for (int mask=1;mask<(1<<(n-1));mask++)
{
int sz=__builtin_popcount(mask); //the size of s
if (sum((1<<n)-1)==sum(mask)) //the combined element is 0
{
for (int i=0;i<n-1;i++)
{
if (mask&(1<<i))
{
dp[mask]=(dp[(mask^(1<<i))]*(MX-sz))%mod;
break;
}
}
}
else //the combined element is non-zero
{
dp[mask]=(fact[MX]*inv[MX-sz])%mod;
for (int i=0;i<n-1;i++)
{
if (mask&(1<<i))
dp[mask]=(dp[mask]-dp[mask^(1<<i)]+mod)%mod;
}
}
}
printf("%d",dp[(1<<(n-1))-1]);
}
, Nice explanation.