比賽鏈接

C. 變化的數組(Easy Version)

題目大意

一個長度為 $n$ 的非負數組 $a$，要求執行 $k$ 次操作，每次操作如下：

• 有 $\frac{1}{2}$ 的概率令 $a_i\leftarrow a_i+(a_i?m)+x,?i\in [1,n]$；
• 另 $\frac{1}{2}$ 的概率保持 $?a_i$ 不變。

求 $\sum _{i=1}^n{a}_i$ 的期望，答案對 $998244353$ 取模。

其中 $?$ 表示按位與，例如 $(10{)}_2?(11{)}_2=(10{)}_2,(01{)}_2?(10{)}_2=0$。

數據范圍

• $1\le n\le 10^6,$
• $1\le m,k\le 5?10^3,$
• $0\le x\le 10^5,$
• $0\le a_i\le 10^9.$

Solution

我們觀察 $a_i$ 的增量 $(a_i?m)+x$，發現除了給定的 $x$，$(a_i?m)$ 只與后 $?{\log }_{2}m?$ 位有關，于是記 $$M=2^{?{\log }_{2}m?+1},$$

這樣一來我們只需要知道 $a_i?(M?1)$ 就能知道 $a_i$ 的增量。

于是我們把每個 $a_i$ 劃分為兩部分，分別是 $?\frac{a_i}{M}?$ 和 $a_i?(M?1)$，我們稱其為高位和低位。

接下來我們就分別求高位和低位的期望 $\mathrm{his}$ 和 $\mathrm{los}$，最終的答案就是 $\mathrm{his}\times \mathrm{M}+\mathrm{los}$。

對于低位來說，我們可以構造一個轉換表 $\mathrm{suf}$，其中 $$\mathrm{suf}[\mathrm{v}]=(\mathrm{v}+(\mathrm{v}?\mathrm{m})+\mathrm{x})?(\mathrm{M}?1),\mathrm{v}\in [0,\mathrm{M}).$$

這樣就可以求出低位和的期望 $\mathrm{los}$。

• 假設有 $j$ 次操作讓 $a_i$ 發生改變，現在已經求出 $j?1$ 次改變時每個低位值的個數，記為 $cn{t}_{j?1}[v]$，其中 $v\in [0,M)$，那么只要對 $?v\in [0,M)$ 做一次 $\mathrm{suf}$ 變換就可以得到新的 $v^{\prime }$ 以及 $cn{t}_j[v^{\prime }]$ 了，具體來說就是 $$\mathrm{cn}{\mathrm{t}}_{\mathrm{j}}[\mathrm{suf}[\mathrm{v}]]:=\mathrm{cn}{\mathrm{t}}_{\mathrm{j}}[\mathrm{suf}[\mathrm{v}]]+\mathrm{cn}{\mathrm{t}}_{\mathrm{j}?1}[\mathrm{v}].$$ 再對每個 $cn{t}_j[v^{\prime }]\times v^{\prime }$ 乘上 $j$ 次改變的概率 $${\left(\frac{1}{2}\right)}^k\left(\genfrac{}{}{0ex}{}{k}{j}\right),$$ 最后求和就是期望。
• 初始值 $cn{t}_0[v]$ 只要對 $?a_i$ 記錄 $a_i?(M?1)$ 的數量即可。

高位就稍微復雜一些。

我們模仿低位，構造一個高位映射表 $\mathrm{pre}$，其中 $$\mathrm{pre}[\mathrm{v}]=?\frac{(\mathrm{v}+(\mathrm{v}?\mathrm{m})+\mathrm{x})}{\mathrm{M}}?,\mathrm{v}\in [0,\mathrm{M}).$$

對于高位來說，我們不用期望的原始公式 $$E(X)=\sum _{i=1}^N{p}_i{X}_i,$$ 而是選擇一個基準 $B$，對其變形得到 $$E(X)=\sum _{i=1}^N{p}_i(X_i?B+B)=B+\sum _{i=1}^N{p}_i(X_i?B).$$ 其中 $(X_i?B)$ 是每個隨機變量取值 $X_i$ 相對于 $B$ 的增量。

在高位上我們選擇的基準 $$B=\sum _{i=1}^n?\frac{(a_i+(a_i?m)+x)}{M}?.$$

接下來就是算高位 增量 的期望了。

我們先給出求和式。

$$\sum _{j=0}^k{\left(\frac{1}{2}\right)}^k\left(\genfrac{}{}{0ex}{}{k}{j}\right)\sum _{i=0}^{j?1}\sum _{v=0}^{M?1}cn{t}_i[v]\times pre[v],$$

上式中 ${\left(\frac{1}{2}\right)}^k\left(\genfrac{}{}{0ex}{}{k}{j}\right)$ 表示對數組 $a$ 做了 $j$ 次改變的概率，后面的兩重循環是求從開始到改變 $j$ 次的增量和。

對于 $j$，之所以我們要遍歷 $i\in [0,j)$，是因為 $cn{t}_i[v]$ 是不斷變化的；而低位不需要這樣遍歷則是因為它不用求增量，可以直接獲得值。

但是這個三重循環的復雜度我們無法接受，所以考慮交換求和次序。

$$\begin{array}{rl}& \sum _{j=0}^k{\left(\frac{1}{2}\right)}^k\left(\genfrac{}{}{0ex}{}{k}{j}\right)\sum _{i=0}^{j?1}\sum _{v=0}^{M?1}cn{t}_i[v]\times pre[v]\\ & =\sum _{i=0}^k\sum _{j=i+1}^k\sum _{v=0}^{M?1}{\left(\frac{1}{2}\right)}^k\left(\genfrac{}{}{0ex}{}{k}{j}\right)\times cn{t}_i[v]\times pre[v]\\ & =\sum _{i=0}^k\sum _{v=0}^{M?1}cn{t}_i[v]\times pre[v]\sum _{j=i+1}^k{\left(\frac{1}{2}\right)}^k\left(\genfrac{}{}{0ex}{}{k}{j}\right)\\ & =\sum _{i=0}^k\sum _{v=0}^{M?1}cn{t}_i[v]\times pre[v]\times s[j].\end{array}$$

其中 $$s[j]=\sum _{j=i+1}^k{\left(\frac{1}{2}\right)}^k\left(\genfrac{}{}{0ex}{}{k}{j}\right).$$

這樣就把復雜度降到 $O(Mk)$ 了。

時間復雜度 $O(mk)$

• 雖然說 $M=2^{?{\log }_{2}m?}+1$，但是量級上最多是 $2m$，$2$ 可以看作常數。

C++ Code

#include <bits/stdc++.h>using u32 = unsigned;
using i64 = long long;
using u64 = unsigned long long;template<class T>
constexpr T power(T a, i64 b) {T res = 1;for (; b; b /= 2, a *= a) {if (b % 2) {res *= a;}}return res;
}
template<int P>
struct MInt {int x;constexpr MInt() : x{} {}constexpr MInt(i64 x) : x{norm(x % getMod())} {}static int Mod;constexpr static int getMod() {if (P > 0) {return P;} else {return Mod;}}constexpr static void setMod(int Mod_) {Mod = Mod_;}constexpr int norm(int x) const {if (x < 0) {x += getMod();}if (x >= getMod()) {x -= getMod();}return x;}constexpr int val() const {return x;}explicit constexpr operator int() const {return x;}constexpr MInt operator-() const {MInt res;res.x = norm(getMod() - x);return res;}constexpr MInt inv() const {assert(x != 0);return power(*this, getMod() - 2);}constexpr MInt &operator*=(MInt rhs) & {x = 1LL * x * rhs.x % getMod();return *this;}constexpr MInt &operator+=(MInt rhs) & {x = norm(x + rhs.x);return *this;}constexpr MInt &operator-=(MInt rhs) & {x = norm(x - rhs.x);return *this;}constexpr MInt &operator/=(MInt rhs) & {return *this *= rhs.inv();}friend constexpr MInt operator*(MInt lhs, MInt rhs) {MInt res = lhs;res *= rhs;return res;}friend constexpr MInt operator+(MInt lhs, MInt rhs) {MInt res = lhs;res += rhs;return res;}friend constexpr MInt operator-(MInt lhs, MInt rhs) {MInt res = lhs;res -= rhs;return res;}friend constexpr MInt operator/(MInt lhs, MInt rhs) {MInt res = lhs;res /= rhs;return res;}friend constexpr std::istream &operator>>(std::istream &is, MInt &a) {i64 v;is >> v;a = MInt(v);return is;}friend constexpr std::ostream &operator<<(std::ostream &os, const MInt &a) {return os << a.val();}friend constexpr bool operator==(MInt lhs, MInt rhs) {return lhs.val() == rhs.val();}friend constexpr bool operator!=(MInt lhs, MInt rhs) {return lhs.val() != rhs.val();}
};template<>
int MInt<0>::Mod = 998244353;template<int V, int P>
constexpr MInt<P> CInv = MInt<P>(V).inv();constexpr int P = 998244353;
using Z = MInt<P>;struct Comb {int n;std::vector<Z> _fac;std::vector<Z> _invfac;std::vector<Z> _inv;Comb() : n{0}, _fac{1}, _invfac{1}, _inv{0} {}Comb(int n) : Comb() {init(n);}void init(int m) {m = std::min(m, Z::getMod() - 1);if (m <= n) return;_fac.resize(m + 1);_invfac.resize(m + 1);_inv.resize(m + 1);for (int i = n + 1; i <= m; i += 1) {_fac[i] = _fac[i - 1] * i;}_invfac[m] = _fac[m].inv();for (int i = m; i > n; i -= 1) {_invfac[i - 1] = _invfac[i] * i;_inv[i] = _invfac[i] * _fac[i - 1];}n = m;}Z fac(int m) {if (m > n) init(2 * m);return _fac[m];}Z invfac(int m) {if (m > n) init(2 * m);return _invfac[m];}Z inv(int m) {if (m > n) init(2 * m);return _inv[m];}Z binom(int n, int m) {if (n < m || m < 0) {return 0;}return fac(n) * invfac(m) * invfac(n - m);}Z Lucas(i64 n, i64 m, int p) {if (n < p and m < p) {return binom(n, m);}return Lucas(n / p, m / p, p) * binom(n % p, m % p);}Z Lucas(i64 n, i64 m) {if (n < Z::getMod() and m < Z::getMod()) {return binom(n, m);}return Lucas(n / Z::getMod(), m / Z::getMod()) * binom(n % Z::getMod(), m % Z::getMod());   }Z perm(int n, int m) {if (n < m or m < 0) {return 0;}return fac(n) * invfac(n - m);}
} comb;template<class T>
std::istream &operator>>(std::istream &is, std::vector<T> &v) {for (auto &x: v) {is >> x;}return is;
}int main() {std::ios::sync_with_stdio(false);std::cin.tie(nullptr);int n, x, m, k;std::cin >> n >> x >> m >> k;std::vector<int> a(n);std::cin >> a;int lm = std::__lg(m) + 1;int M = 1 << lm;std::vector<int> pre(M);std::vector<int> suf(M);for (int i = 0; i < M; i++) {int v = i + (i & m) + x;pre[i] = v >> lm;suf[i] = v & (M - 1);}Z hi0 = 0;std::vector<int> cnt(M);for (int ai: a) {cnt[ai & (M - 1)]++;hi0 += ai >> lm;}std::vector<Z> binom(k + 1);for (int i = 0; i <= k; i++) {binom[i] = comb.binom(k, i) / power(Z(2), k);}std::vector<Z> s(k + 2);for (int i = k; i >= 0; i--) {s[i] = s[i + 1] + binom[i];}Z los = 0;Z his = hi0;for (int i = 0; i <= k; i++) {Z lo = 0;Z hi = 0;for (int j = 0; j < M; j++) {lo += Z(cnt[j]) * j;hi += Z(cnt[j]) * pre[j];}los += binom[i] * lo;his += s[i + 1] * hi;std::vector<int> ncnt(M);for (int j = 0; j < M; j++) {ncnt[suf[j]] += cnt[j];}cnt = std::move(ncnt);}Z ans = his * M + los;std::cout << ans << "\n";return 0;
}