CF1967C Fenwick Tree

树状数组，计数

非常好玩的题，考察对树状数组本质形态的理解。

Link：https://codeforces.com/contest/1967/problem/C

简述

对于数列 \(a\)，定义数列 \(f(a)\) 为对数列 \(a\) 建立的树状数组。即有：

\[f(a)_k = \sum_{i = k - \operatorname{lowbit}(k) + 1}^k a_i \]

给定 \(T\) 组数据，每组数据给定长度为 \(n\) 的数列 \(b\)，参数 \(k\)。请构造长度为 \(n\) 的数列 \(a\)，使得数列 \(a\) 在模 998244353 意义下满足 \(b = f^k(a)\)。可以证明对于任意数列 \(b\) 和参数 \(k\) 总存在一种合法的 \(a\) 的构造方案。
\(1\le T\le 10^4\)，\(1\le n\le 2\cdot 10^5\)，\(1\le k\le 10^9\)，\(0\le b_i< 998244353\)。
3S，256MB。

分析

树状数组本质上是二进制优化的前缀和，由题目给定的式子可以看出来，每个树状数组的位置 \(s_k\) 实际上是以 \(k\) 为右端点的，长度为 \(\operatorname{lowbit}(k)\) 的区间的和。

而本题对给定数组做了 \(k\) 次建立树状数组的操作，实际上即对所有树状数组区间做了 \(k\) 维二进制优化的前缀和。众所周知做高维前缀和时，每个位置对之后的位置的贡献的系数由下表所示，由组合意义易知位于 \((x, y)\) 的系数的值即 \({{x + y}\choose {y}}\)：

\(a\)	1	1	1	1
\(f^1\)	1	2	3	4
\(f^2\)	1	3	6	10
\(f^3\)	1	4	10	20
\(f^4\)	1	5	15	35

记需要还原出来的数列为 \(a\)，考虑对于每个位置 \(i\) 会对哪些位置 \(p\) 的高维前缀和 \(f^k(p)\) 有贡献。由建树状数组的过程可知，\(p\) 可以通过树状数组插入操作进行枚举，易知此类位置只有 \(O(\log_2 n)\) 数量级。

根据高维前缀和的式子，手玩下系数矩阵，发现若位置 \(p\) 是在树状数组插入操作时第 \(j(j\ge 0)\) 个被枚举到的位置，则 \(a_i\) 对 \(f^k(p)\) 的贡献的系数即为：

\[{{k+j-1}\choose{j}} = \frac{(k+j-1)}{j}{{k+j-2}\choose{j-1}} \]

于是考虑顺序枚举所有位置并进行还原，还原第 \(i\) 个位置的值 \(a_i\) 后，考虑枚举它有贡献的位置并减去 \(a_i\) 的贡献即可。

总时间复杂度 \(O(n\log n)\) 级别，代码非常好写！

哎呦我草太优美啦树状数组！

代码

//
/*
By:Luckyblock
*/
#include <bits/stdc++.h>
#define LL long long
#define lowbit(x) ((x)&(-x))
const int kN = 2e5 + 10;
const LL p = 998244353;
//=============================================================
int n, k;
LL b[kN], a[kN], ans[kN], inv[kN];
//=============================================================
//=============================================================
int main() {
  //freopen("1.txt", "r", stdin);
  std::ios::sync_with_stdio(0), std::cin.tie(0);
  inv[1] = 1;
  for (int i = 2; i <= 2e5; ++ i) {
    inv[i] = 1ll * (p - p / i + p) % p * inv[p % i] % p;
  }

  int T; std::cin >> T;
  while (T --) {
    std::cin >> n >> k;
    for (int i = 1; i <= n; ++ i) std::cin >> b[i], a[i] = b[i];
    for (int i = 1; i <= n; ++ i) {
      ans[i] = a[i];
      LL c = 1, d = 0;
      for (int j = i; j <= n; j += lowbit(j)) {
        a[j] = (a[j] - c * ans[i] % p + p) % p;
        c = c * (k + d) % p * inv[d + 1] % p;
        ++ d;
      }
    }
    for (int i = 1; i <= n; ++ i) std::cout << ans[i] << " ";
    std::cout << "\n";
  }

  return 0;
}
/*
打表看对位置 8 有贡献的位置的的系数矩阵的代码：

int a[15][15][15] = {0};
  for (int i = 1; i <= 8; ++ i) a[0][i][i] = 1;
  
  for (int i = 1; i <= 10; ++ i) {
    for (int j = 1; j <= 8; ++ j) {
      if (j % 2 == 1) {
        a[i][j][j] = 1;
        continue;
      }
      if (j == 2 || j == 6) {
        a[i][j][j] = 1, a[i][j][j - 1] = i;
      }
      if (j == 4) {
        for (int k = j - 3; k <= j; ++ k) {
          for (int l = 1; l <= 8; ++ l) {
            a[i][j][l] += a[i - 1][k][l];
          }
        }
      }
      if (j == 8) {
        for (int k = j - 7; k <= j; ++ k) {
          for (int l = 1; l <= 8; ++ l) {
            a[i][j][l] += a[i - 1][k][l];
          }
        }
      }
    }
    std::cout << i << ": ";
    for (int j = 1; j <= 8; ++ j) printf("%5d", a[i][8][j]);
    std::cout << "\n";
  }
*/