[ARC200E] popcount<=2

题目大意

给定 \(N,M\)，你需要求有多少个长为 \(N\) 的序列 \(A\)，满足下列条件：

• \(0\leq A_i\leq 2^M(1\leq i\leq N)\)
• \(\text{popcount}(A_i\oplus A_j)\leq 2(1\leq i\leq j\leq N)\)

\(T\) 组测试数据

• \(1\leq T\leq 2\times 10^5\)
• \(2\leq N,M<998244353\)

解题思路

特别的，下面每个部分开头将介绍怎么想到这么做。

[0] 化简限制

\(A_i\) 取值可能太多，不方便处理。考虑化简。

容易发现令 \(A\leftarrow (A_1\oplus A_1,A_2\oplus A_1,\dots,A_N\oplus A_1)\) 后合法性不变。记得算答案的时候方案数乘以 \(2^M\)，表示 \(A_1\) 的取值种类数。

此时的 \(A\) 满足：

• \(A_1=0\)
• \(\text{popcount}(A_i)\leq 2\)

特别的，为了叙述方便，我们令 \(S=\{A_i|\text{popcount}(A_i)=2\}\)。

[1] \(\text{popcount}(A_i)\leq 1\)，即 \(|S|=0\)

假若存在 \(\text{popcount}(A_i)=2\) 的，显然是不好处理的。所以先从简单情况考虑。

对于每个 \(i\)，有 \(\text{popcount}(A_i)\leq 1\)。此时有 \(\text{popcount}(A_i\oplus A_j)\leq\text{popcount}(A_i)+\text{popcount}(A_j)\leq 2\) 成立。

因此对于 \(A\) 的唯一限制是 \(\text{popcount}(A_i) \leq 1\)。满足此条件的序列 \(A\) 共有 \((M+1)^{N-1}\) 种。

[2] \(|S|=1\)

感觉上，可能会被用到的/需要特殊计数的数位并不多，所以考虑对 \(|S|\) 的大小进行分类讨论。容易发现 \(|S|=1\) 时，\(|S|\) 内的元素不存在两两间的限制，需要单独处理。

记 \(S=\{2^x+2^y\}\)，那么：

• \(x,y\) 两数共有 \(\binom{M}{2}\) 种选法。

  也就是，\(x\) 有 \(\binom{M}{2}\) 种取值。

• \(\text{popcount}(A_i)=0\)，有 \(1\) 种可能值：\(A_i=0\)

• \(\text{popcount}(A_i)=1\)，有 \(2\) 种可能值：\(A_i=2^x,2^y\)

• \(\text{popcount}(A_i)=2\)，有 \(1\) 种可能值：\(A_i=2^x+2^y\)

• \(A_i=2^x+2^y\) 出现至少一次。

难以直接计算的是“存在 \(A_i=2^x+2^y\)”这一限制。实际上，在确定了 \(x,y\) 之后，拿总方案数 \(4^{N-1}\) 减去不合法的方案数 \(3^{N-1}\) 即可。

因此，共有 \(\binom{M}{2}(4^{M-1}-3^{M-1})\) 种方案。

[3] \(|S|=2\)

此时，对于任意 \(|S|\) 中的不同元素 \(a,b\)，满足两者的二进制位集合交集非空。归纳容易发现，我们可以将情况分为两种。

[3.1] \(S\) 中所有元素的按位与结果为 \(0\)

记 \(S\) 中的第一个元素为 \(2^x+2^y\)，第二个元素为 \(2^x+2^z\)。那么根据两两交集非空的限制，第三者的可能是 \(2^y+2^z\)。假若如此，那么此时明显不能有更多的元素了。

记当前 \(S=\{2^x+2^y,2^x+2^z,2^y+2^z\}\)。注意到，对于任意 \(a,b\in S\)，\(\text{popcount}(A_i)\leq 2\)。并且，此时不会存在 \(\text{popcount}(A_i)=1\) 的情况

那么有：

• 确定 \(x,y,z\) 共有 \(\binom{M}{3}\) 种方式。
• \(\text{popcount}(A_i)=0\)，有 \(1\) 种可能值：\(A_i=0\)
• \(\text{popcount}(A_i)=2\)，有 \(3\) 种可能值：\(A_i=2^x+2^y,2^x+2^z,2^y+2^z\)
• 对于 \(S\) 中的元素，每种都应当出现至少一次。

与 \(|S|=1\) 同理，在确定了 \(x,y,z\) 之后，可以容斥得到满足最后一条限制的方案数： \(\binom{3}{3}4^{N-1}-\binom{3}{2}3^{N-1}+\binom{3}{1}2^{N-1}-\binom{3}{0}1^{N-1}\)。

化简之后，序列 \(A\) 的选择方案共有 \(\binom{M}{3}(4^{N-1}-3\times 3^{N-1}+3\times 2^{N-1}-1)\) 种。

[3.2] \(S\) 中所有元素的按位与非 \(0\)

根据上述的，我们假设 \(S\) 中前两个元素分别为 \(2^x+2^y,2^x+2^z\)。那么第三者还可以是 \(2^x+2^t\) 这样的。甚至，\(S\) 中还可以存在更多的形如 \(2^x+2^{\dots}\) 这样的数。

考虑 \(S\) 中所有元素共享某个二进制位（记为 \(2^x\)）的情况。此时，如果 \(a,b\in S\)，满足 \(\text{popcount}(A_i)\leq 2\)。对于 \(\text{popcount}(A_i)=1\) 的情况，明显只有 \(A_i=2^x\) 这一种选择。

因此：

• \(x\) 有 \(M\) 种选择方式
• \(\text{popcount}(A_i)=1\)，共有 \(1\) 种取值：\(A_i=2^x\)
• \(\text{popcount}(A_i)=0\)，共有 \(1\) 种取值：\(A_i=0\)
• \(\text{popcount}(A_i) = 2\)，共有 \(M-1\) 种取值：\(A_i=2^x+2^{\dots}\)
• 必须有至少两个 \(A_i\) 满足 \(\text{popcount}(A_i) = 2\)

因此，满足条件的序列 \(A\) 的数量为 \(M\left((M+1)^{N-1}-(M-1)(3^{N-1}-2^{N-1}) - 2^{N-1}\right)\)。

其中，\((M+1)^{N-1}\) 是总方案数，\((M-1)(3^{N-1}-2^{N-1})\) 是只有一个 \(\text{popcount}(A_i)=2\) 的方案数，\(2^{N-1}\) 是不存在的方案数。

---

综上，我们可以在 \(O(\log N)\) 的复杂度内处理一组测试数据。总时间复杂度为 \(O(T\log N)\)。

代码实现

点击查看代码

#include <bits/stdc++.h>
#include <atcoder/all>
#define FL(i, a, b) for (int i = (a); i <= (b); ++i)
#define FR(i, a, b) for (int i = (a); i >= (b); --i)
using namespace std;
using mint = atcoder::modint998244353;
void Solve() {
	int n; mint m;
	scanf("%d %u", &n, &m);
	mint c2 = m * (m - 1) / 2;
	mint c3 = m * (m - 1) * (m - 2) / 6;
	mint p2 = mint(2).pow(n - 1);
	mint p3 = mint(3).pow(n - 1);
	mint p4 = mint(4).pow(n - 1);
	mint pw = (m + 1).pow(n - 1);
	mint ans = pw;
	ans += c2 * (p4 - p3);
	ans += c3 * (p4 - 3 * p3 + 3 * p2 - 1);
	ans += m * (pw - p2 - (m - 1) * (p3 - p2));
	ans *= mint(2).pow(m.val());
	printf("%u\n", ans);
}
int main() {
	int T;
	scanf("%d", &T);
	while (T--) {
		Solve();
	}
	return 0;
}