PREFSUFANDOR - Editorial

PROBLEM LINK:

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

Author: Jeevan Jyot Singh
Testers: Abhinav Sharma, Venkata Nikhil Medam
Editorialist: Nishank Suresh

DIFFICULTY:

2854

PREREQUISITES:

Observation, basic combinatorics

PROBLEM:

JJ has an array A of length N, where each element lies between 0 and 2^K - 1. In one move, he can pick two integers x and i such that 0 \leq x \lt 2^K and 1 \leq i \leq N, and:

  • Set A_j := A_j \And x for the prefix ending at i, or
  • Set A_j := A_j \mid x for the suffix starting at i

How many different arrays can he create by performing this operation several times?

EXPLANATION:

First, note that the problem can be solved for each bit independently since there is no limit on the number of operations. So, if we are able to solve the problem for K = 1, we can apply this solution to every bit separately and multiply all the answers together to obtain the final answer.

Now, we just need to solve the problem for K = 1, i.e, the array contains only 0-s and 1-s. When this is the case, our operations reduce to:

  • Setting some prefix of the array to 0 (the first operation, with x = 0)
  • Setting some suffix of the array to 1 (the second operation, with x = 1)

Note that performing the first operation with x = 1 or the second operation with x = 0 don’t change the array at all, so we can safely ignore them.

This gives us the following observations:

  • It is enough to perform at most one operation of each type
  • The prefix chosen for the first operation and the suffix chosen for the second can be chosen in such a way that they don’t intersect.
  • If an operation is performed on the prefix of length i, and A_i = 0 initially, we can achieve the same result by performing the operation on the prefix of length i-1 instead. So, it is enough to consider the prefix operation only for those i such that A_i = 1.
  • Similarly, it is enough to consider the suffix operation only for those i such that A_i = 0.

Putting all this together, we can see that the problem simply reduces to counting the number of \texttt{10}-subsequences of A (do you see why?), which can be done easily in \mathcal{O}(N).

Don’t forget to account for the case where no prefix operation and/or no suffix operation is performed!

As mentioned at the start, the final answer is obtained by solving this subproblem for each of the K bits, and multiplying all the answers together.

TIME COMPLEXITY:

\mathcal{O}(N \cdot K) per test case.

CODE:

Setter (C++)
#include <bits/stdc++.h>
using namespace std;

#define int long long
#define endl "\n"
#define sz(w) (int)(w.size())

const int mod = 998244353;

void solve()
{
    int n, k; cin >> n >> k;
    vector<int> a(n);
    for(int &x: a) cin >> x;
    int ans = 1;
    for(int i = 0; i < k; i++)
    {
        vector<int> bit(n);
        for(int j = 0; j < n; j++)
            bit[j] = (a[j] >> i) & 1;
        int cnt0 = count(bit.begin(), bit.end(), 0);
        int res = cnt0 + 1;
        for(int j = 0; j < n; j++)
        {
            if(bit[j] == 0)
                cnt0--;
            else
                res += cnt0 + 1;
        }
        ans = ans * (res % mod) % mod;
    }    
    cout << ans << endl;
}

int32_t main()
{
    ios::sync_with_stdio(0); 
    cin.tie(0);
    int T; cin >> T;
    while(T--)
        solve();
    return 0;
}
Tester (nikhil_medam, C++)
// Tester: Nikhil_Medam
#include <bits/stdc++.h>
using namespace std;
#define endl "\n"
#define int long long
const int N = 1e5 + 5;
const int mod = 998244353;

int t, n, k, a[N];
int32_t main() {
    cin >> t;
    while(t--) {
        cin >> n >> k;
        for(int i = 0; i < n; i++) {
            cin >> a[i];
        }
        int ans = 1;
        for(int i = 0; i < k; i++) {
            int bit[n], cnt_0;
            for(int j = 0; j < n; j++) {
                bit[j] = (a[j] >> i) & 1;
                cnt_0 += 1 - bit[j];
            }
            int cur_ans = cnt_0 + 1;
            for(int j = 0; j < n; j++) {
                if(bit[j] == 0) {
                    cnt_0--;
                }
                else {
                    cur_ans += cnt_0 + 1;
                }
            }
            ans = (ans * (cur_ans % mod)) % mod;
        }
        cout << ans << endl;
    }
	return 0;
}
Editorialist (Python)
mod = 998244353
def solve(a):
	z = a.count(0)
	ret = 0
	for x in a:
		if x == 1:
			ret += z
		else:
			z -= 1
	return ret
for _ in range(int(input())):
	n, k = map(int, input().split())
	a = list(map(int, input().split()))
	a.append(0)
	a.insert(0, (1<<k) - 1)
	ans = 1
	for bit in range(k):
		b = []
		for x in a:
			b.append((x >> bit)&1)
		ans *= solve(b)
		ans %= mod
	print(ans)

Can you please prove why?
It doesn’t seem obvious at the first glance.

The four bullet points mentioned just above it contain all the steps needed for the proof. I still encourage you to try proving it yourself using those, read the spoiler below if you cannot.

Proof

Suppose you perform only one prefix operation, and no suffix operation. How many strings are possible?
Well, we know that it’s optimal to perform the prefix operation only on a length i such that A_i = 1. Performing this operation sets everything in this prefix to 0, so the number of distinct operations is simply the number of 1-s in the string.

Similarly, the number of distinct suffix operations is the number of 0-s in the string.

Now, what if you want to combine a prefix operation and a suffix operation?
As mentioned, the operations will not overlap (if they do, you can make one of them shorter without affecting the final result).

So, say you decide to perform a prefix operation on length i. How many suffix operations can be matched with this prefix operation? From the reasoning above, The answer is simply the number of 0-s to the right of i.

You can see this as picking a 10-subsequence, performing a prefix operation upto the 1, and performing a suffix operation from the 0. Clearly, all of these are going to be distinct, and are also the only ones possible.