CF1967C Fenwick Tree

Luckyblock發表於2024-09-11

樹狀陣列,計數

非常好玩的題,考察對樹狀陣列本質形態的理解。

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";
  }
*/