LCP 05. 发 LeetCoin

LCP 05. 发 LeetCoin

树状数组

dfs得到begin和end数组也很灵魂。

用树状数组的经典问题：
对一个数组中，区间加，区间查询数字和

typedef long long LL;
const int N = 50010;
const int MOD = 1e9 + 7;
int n;
LL tr1[N], tr2[N];

int lowbit(int x) { return x & -x; }

void add(LL tr[], int x, LL v) {
  for (int i = x; i <= n; i += lowbit(i)) {
    tr[i] += v;
  }
}

LL get(LL tr[], int x) {
  LL res = 0;
  for (int i = x; i; i -= lowbit(i)) res += tr[i];
  return res;
}

LL getsum(int x) { return 1ll * get(tr1, x) * (x + 1) - get(tr2, x); }

class Solution {
 public:
  int begin[N], end[N], arr[N];
  int idx = 1;

  void init() {
    memset(tr1, 0, sizeof tr1);
    memset(tr2, 0, sizeof tr2);
  }

  vector<vector<int>> g;

  void dfs(int u) {
    begin[u] = idx++;
    for (auto nxt : g[u]) dfs(nxt);
    end[u] = idx - 1;
  }

  vector<int> bonus(int _n, vector<vector<int>>& leadership,
                    vector<vector<int>>& operations) {
    init();
    n = _n;
    g.resize(n + 1, {});

    for (auto p : leadership) g[p[0]].push_back(p[1]);
    dfs(1);

    vector<int> res;
    for (auto p : operations) {
      if (p[0] == 1) {
        int l = begin[p[1]], r = begin[p[1]], v = p[2];
        add(tr1, l, v), add(tr1, r + 1, -v);
        add(tr2, l, 1ll * l * v), add(tr2, r + 1, -1ll * (r + 1) * v);
      } else if (p[0] == 2) {
        int l = begin[p[1]], r = end[p[1]], v = p[2];
        add(tr1, l, v), add(tr1, r + 1, -v);
        add(tr2, l, 1ll * l * v), add(tr2, r + 1, -1ll * (r + 1) * v);
      } else {
        int l = begin[p[1]], r = end[p[1]];
        LL ans = getsum(r) - getsum(l - 1);
        res.push_back(ans % MOD);
      }
    }

    return res;
  }
};

线段树（懒标记）

还没学会