Problem p01771 (C++) — buggy submission, verdict: Time Limit Exceeded

#include <bits/stdc++.h>
typedef long long int ll;
typedef unsigned long long int ull;
#define BIG_NUM 2000000000
#define HUGE_NUM 99999999999999999
#define MOD 1000000007
#define EPS 0.000000001
using namespace std;

#define NUM 150005

struct Data {
  int parent, head, depth, tail, size;
};

int first_N, N;
int root, global_index;
;
ll uniformity[8 * NUM], partial[8 * NUM];
int LEFT[NUM], RIGHT[NUM];
int MAX_LOG_V = 18;
int POW[19];
int parent[18][NUM], depth[NUM], chain_max_depth[NUM];
int index_table[NUM];
Data data[NUM];
vector<int> G[NUM], ADJ[NUM];

void init(int first_N) {
  while (N < first_N)
    N *= 2;
}

int makeTree(int node_id, int parent, int depth) {

  data[node_id].depth = depth;
  data[node_id].parent = parent;
  data[node_id].size = 1; // 部分木のノード数(自分&子孫)

  int child;

  for (int i = 0; i < G[node_id].size(); i++) {

    child = G[node_id][i];
    if (child == parent)
      continue;

    data[node_id].size += makeTree(child, node_id, depth + 1);
  }

  return data[node_id].size;
}

void HLD(int node_id, int head, int parent, int d) {

  index_table[node_id] =
      global_index++; // ★★根に近いものほど、番号が小さい(★左にある★)
  data[node_id].head = head;

  if (chain_max_depth[head] < d) {
    chain_max_depth[head] = d;
    data[head].tail = node_id; // 鎖の末尾を設定
  }

  if (G[node_id].size() == 0)
    return;

  // 子の中で、最大の重み(部分木のノード数)を持つものを探す
  int max_size = -BIG_NUM;
  int max_node_id = -1, child;

  for (int i = 0; i < G[node_id].size(); i++) {

    child = G[node_id][i];
    if (child == parent)
      continue;

    if (max_size < data[child].size) {

      max_size = data[child].size;
      max_node_id = child;
    }
  }

  if (max_node_id == -1)
    return;

  // 最大の重みを持つ子へ(headは変わらず)
  HLD(max_node_id, head, node_id, d + 1);

  // その他の子を処理(新しいラインを作る)
  for (int i = 0; i < G[node_id].size(); i++) {

    child = G[node_id][i];
    if (child == parent || child == max_node_id)
      continue;

    HLD(child, child, node_id, d + 1);
    ADJ[head].push_back(child);
  }
}

void add_and_updateNodes(int left, int right, ll value, int node_id,
                         int node_left, int node_right) {

  if (right < node_left || left > node_right)
    return; // 範囲に重なりがなければ何もしない
  else if (
      left <= node_left &&
      right >=
          node_right) { // このノードがカバーしている区間が、集計区間の部分区間になっている場合
    uniformity[node_id] += value; // 「一様に加えられた値」を加算する
  } else {                        // 部分的な重なりがある場合
    partial[node_id] +=
        (min(right, node_right) - max(left, node_left) + 1) * value;
    add_and_updateNodes(left, right, value, 2 * node_id + 1, node_left,
                        (node_left + node_right) / 2); // 左の子に処理を投げる
    add_and_updateNodes(left, right, value, 2 * node_id + 2,
                        (node_left + node_right) / 2 + 1,
                        node_right); // 右の子に処理を投げる
  }
}

// 再帰的に和を計算し、返却する関数
ll getSum(int left, int right, int node_id, int node_left, int node_right) {
  if (right < node_left || left > node_right)
    return 0;
  else if (
      left <= node_left &&
      right >=
          node_right) { // このノードがカバーしている区間が、集計区間の部分区間になっている場合
    return (node_right - node_left + 1) * uniformity[node_id] +
           partial[node_id];
  } else { // 部分的に一致している場合
    ll sum = (min(right, node_right) - max(left, node_left) + 1) *
             uniformity[node_id];
    sum += getSum(left, right, 2 * node_id + 1, node_left,
                  (node_left + node_right) / 2); // 左の子の和を加算
    sum +=
        getSum(left, right, 2 * node_id + 2, (node_left + node_right) / 2 + 1,
               node_right); // 右の子の和を加算
    return sum;
  }
}

// node_idの子孫の辺の重みをvalueだけ増やす
void add(int node_id, ll value) {

  if (data[node_id].size >= 2) {
    add_and_updateNodes(
        index_table[node_id] + 1, index_table[node_id] + data[node_id].size - 1,
        value, 0, 0,
        N - 1); // ★★次の番号～部分木のサイズ-1の番号が、ちょうど子孫になっている★★
  }
}

// node_id～rootへのコストを求める
ll get_weight(int node_id) {

  ll ret = 0;
  int head_node;

  while (true) {

    head_node = data[node_id].head;

    // node_id～同じラインのヘッドのコストを求める(セグ木演算)
    ret += getSum(index_table[head_node], index_table[node_id], 0, 0,
                  N - 1); // 一様に加えられた値

    if (head_node == root) {
      break; // rootまで到達したらbreak
    }

    node_id = data[head_node].parent;

    if (node_id == root) {
      break;
    }
  }

  return ret;
}

// parentとdepthを再帰的に設定
void dfs(int node_id, int parent_id, int d) {
  parent[0][node_id] = parent_id;
  depth[node_id] = d;
  for (int i = 0; i < G[node_id].size(); i++) {
    if (G[node_id][i] != parent_id)
      dfs(G[node_id][i], node_id, d + 1);
  }
}

// 初期化
void makeLA() {
  // parent[0]とdepthを初期化する
  dfs(root, -1, 0);
  // parentを初期化する
  for (int k = 0; k + 1 < MAX_LOG_V; k++) {
    for (int node_id = 0; node_id < first_N; node_id++) { // ★★ここはfirst_N★★
      if (parent[k][node_id] < 0)
        parent[k + 1][node_id] =
            -1; // node_idの2^k上のノードがルートノードより上なら、2^(k+1)上も同様とする
      else {
        parent[k + 1][node_id] = parent[k][parent[k][node_id]];
      }
    }
  }
}

// leftとrightのLCAを求める
int LCA(int left, int right) {

  if (depth[left] > depth[right])
    swap(left, right);

  for (int k = MAX_LOG_V - 1; k >= 0; k--) {
    if (depth[right] - depth[left] >= POW[k]) {
      right = parent[k][right];
    }
  }

  if (left == right)
    return left;

  for (int k = MAX_LOG_V - 1; k >= 0; k--) {
    if (parent[k][left] != parent[k][right]) {
      left = parent[k][left];
      right = parent[k][right];
    }
  }
  return parent[0][left];
}

ll calc(int from, int to) {

  int lca = LCA(from, to);
  return get_weight(from) + get_weight(to) - 2 * get_weight(lca);
}

int main() {

  POW[0] = 1;
  for (int i = 1; i <= 18; i++) {
    POW[i] = POW[i - 1] * 2;
  }

  int num_query;
  scanf("%d %d", &first_N, &num_query);

  int from, to;

  for (int loop = 0; loop < first_N - 1; loop++) {

    scanf("%d %d", &from, &to);
    G[from].push_back(to);
    G[to].push_back(from);
  }

  for (int i = 0; i < first_N; i++) {

    chain_max_depth[i] = -1; // 鎖の中の最大の深さ
  }

  root = 0;
  makeLA();
  makeTree(root, -1, 0);
  global_index = 0;
  HLD(root, root, -1, 0); // 0スタートでdfs木を作ったので、0がrootとなる

  // dfs木をセグメント木に変換する
  N = 1;
  init(first_N);

  for (int i = 0; i < N; i++) {

    uniformity[i] = 0;
    partial[i] = 0;
  }

  int command, node_id;
  ll add_weight;

  for (int loop = 0; loop < num_query; loop++) {

    scanf("%d", &command);

    if (command == 0) { // 重みの和を計算

      scanf("%d %d", &from, &to);
      printf("%lld\\n", calc(from, to));

    } else {

      scanf("%d %lld", &node_id, &add_weight);
      add(node_id, add_weight);
    }
  }

  return 0;
}


Error: TLE

Fixed submission:

#include <bits/stdc++.h>
typedef long long int ll;
typedef unsigned long long int ull;
#define BIG_NUM 2000000000
#define HUGE_NUM 99999999999999999
#define MOD 1000000007
#define EPS 0.000000001
using namespace std;

#define NUM 150005

struct Data {
  int parent, head, depth, tail, size;
  ll sum_add_weight;
};

int first_N, N;
int root, global_index;
;
ll POOL[8 * NUM], uniformity[8 * NUM], partial[8 * NUM];
int LEFT[NUM], RIGHT[NUM];
int MAX_LOG_V = 18;
int POW[19];
int parent[18][NUM], depth[NUM], chain_max_depth[NUM];
int index_table[NUM];
Data data[NUM];
vector<int> G[NUM], ADJ[NUM];

void init(int first_N) {
  while (N < first_N)
    N *= 2;
}

int makeTree(int node_id, int parent, int depth) {

  data[node_id].depth = depth;
  data[node_id].parent = parent;
  data[node_id].sum_add_weight = 0; // 鎖間の重みを保持
  data[node_id].size = 1;           // 部分木のノード数(自分&子孫)

  int child;

  for (int i = 0; i < G[node_id].size(); i++) {

    child = G[node_id][i];
    if (child == parent)
      continue;

    data[node_id].size += makeTree(child, node_id, depth + 1);
  }

  return data[node_id].size;
}

void HLD(int node_id, int head, int parent, int d) {

  index_table[node_id] =
      global_index++; // ★★根に近いものほど、番号が小さい(★左にある★)
  data[node_id].head = head;

  if (chain_max_depth[head] < d) {
    chain_max_depth[head] = d;
    data[head].tail = node_id; // 鎖の末尾を設定
  }

  if (G[node_id].size() == 0)
    return;

  // 子の中で、最大の重み(部分木のノード数)を持つものを探す
  int max_size = -BIG_NUM;
  int max_node_id = -1, child;

  for (int i = 0; i < G[node_id].size(); i++) {

    child = G[node_id][i];
    if (child == parent)
      continue;

    if (max_size < data[child].size) {

      max_size = data[child].size;
      max_node_id = child;
    }
  }

  if (max_node_id == -1)
    return;

  // 最大の重みを持つ子へ(headは変わらず)
  HLD(max_node_id, head, node_id, d + 1);

  // その他の子を処理(新しいラインを作る)
  for (int i = 0; i < G[node_id].size(); i++) {

    child = G[node_id][i];
    if (child == parent || child == max_node_id)
      continue;

    HLD(child, child, node_id, d + 1);
    ADJ[head].push_back(child);
  }
}

void add_and_updateNodes(int left, int right, ll value, int node_id,
                         int node_left, int node_right) {

  if (right < node_left || left > node_right)
    return; // 範囲に重なりがなければ何もしない
  else if (
      left <= node_left &&
      right >=
          node_right) { // このノードがカバーしている区間が、集計区間の部分区間になっている場合
    uniformity[node_id] += value; // 「一様に加えられた値」を加算する
  } else {                        // 部分的な重なりがある場合
    partial[node_id] +=
        (min(right, node_right) - max(left, node_left) + 1) * value;
    add_and_updateNodes(left, right, value, 2 * node_id + 1, node_left,
                        (node_left + node_right) / 2); // 左の子に処理を投げる
    add_and_updateNodes(left, right, value, 2 * node_id + 2,
                        (node_left + node_right) / 2 + 1,
                        node_right); // 右の子に処理を投げる
  }
}

// 再帰的に和を計算し、返却する関数
ll getSum(int left, int right, int node_id, int node_left, int node_right) {
  if (right < node_left || left > node_right)
    return 0;
  else if (
      left <= node_left &&
      right >=
          node_right) { // このノードがカバーしている区間が、集計区間の部分区間になっている場合
    return (node_right - node_left + 1) * uniformity[node_id] +
           partial[node_id];
  } else { // 部分的に一致している場合
    ll sum = (min(right, node_right) - max(left, node_left) + 1) *
             uniformity[node_id];
    sum += getSum(left, right, 2 * node_id + 1, node_left,
                  (node_left + node_right) / 2); // 左の子の和を加算
    sum +=
        getSum(left, right, 2 * node_id + 2, (node_left + node_right) / 2 + 1,
               node_right); // 右の子の和を加算
    return sum;
  }
}

// node_idの子孫の辺の重みをvalueだけ増やす
void add(int node_id, ll value) {

  int head_node, tail_node;

  head_node = data[node_id].head;

  // 区間(鎖)に値を足す(セグ木演算)
  if (head_node != node_id) {
    tail_node = data[head_node].tail;
    add_and_updateNodes(index_table[node_id], index_table[tail_node], value, 0,
                        0, N - 1);
    add_and_updateNodes(index_table[node_id], index_table[node_id], -value, 0,
                        0, N - 1); // ★★注意★★
  } else {

    POOL[head_node] += value;
  }

  int child; // 鎖にくっついている子(★★自分の子孫方向★★)

  for (int i = 0; i < ADJ[data[node_id].head].size(); i++) {

    child = ADJ[data[node_id].head][i];
    if (depth[child] <= depth[node_id])
      continue;

    data[child].sum_add_weight += value; // 鎖と鎖の間
    add(child, value);
  }
}

// node_id～rootへのコストを求める
ll get_weight(int node_id) {

  ll ret = 0;
  int head_node;

  while (true) {

    head_node = data[node_id].head;

    ret += POOL[head_node] * (depth[node_id] - depth[head_node]);

    // node_id～同じラインのヘッドのコストを求める(セグ木演算)
    if (head_node != node_id) {

      ret += getSum(index_table[head_node], index_table[node_id], 0, 0,
                    N - 1); // 一様に加えられた値
    }

    if (head_node == root) {
      break; // rootまで到達したらbreak
    }

    ret += data[head_node].sum_add_weight; // 鎖間の重み

    node_id = data[head_node].parent;

    if (node_id == root) {
      break;
    }
  }

  return ret;
}

// parentとdepthを再帰的に設定
void dfs(int node_id, int parent_id, int d) {
  parent[0][node_id] = parent_id;
  depth[node_id] = d;
  for (int i = 0; i < G[node_id].size(); i++) {
    if (G[node_id][i] != parent_id)
      dfs(G[node_id][i], node_id, d + 1);
  }
}

// 初期化
void makeLA() {
  // parent[0]とdepthを初期化する
  dfs(root, -1, 0);
  // parentを初期化する
  for (int k = 0; k + 1 < MAX_LOG_V; k++) {
    for (int node_id = 0; node_id < first_N; node_id++) { // ★★ここはfirst_N★★
      if (parent[k][node_id] < 0)
        parent[k + 1][node_id] =
            -1; // node_idの2^k上のノードがルートノードより上なら、2^(k+1)上も同様とする
      else {
        parent[k + 1][node_id] = parent[k][parent[k][node_id]];
      }
    }
  }
}

// leftとrightのLCAを求める
int LCA(int left, int right) {

  if (depth[left] > depth[right])
    swap(left, right);

  for (int k = MAX_LOG_V - 1; k >= 0; k--) {
    if (depth[right] - depth[left] >= POW[k]) {
      right = parent[k][right];
    }
  }

  if (left == right)
    return left;

  for (int k = MAX_LOG_V - 1; k >= 0; k--) {
    if (parent[k][left] != parent[k][right]) {
      left = parent[k][left];
      right = parent[k][right];
    }
  }
  return parent[0][left];
}

ll calc(int from, int to) {

  int lca = LCA(from, to);
  return get_weight(from) + get_weight(to) - 2 * get_weight(lca);
}

int main() {

  POW[0] = 1;
  for (int i = 1; i <= 18; i++) {
    POW[i] = POW[i - 1] * 2;
  }

  int num_query;
  scanf("%d %d", &first_N, &num_query);

  int from, to;

  for (int loop = 0; loop < first_N - 1; loop++) {

    scanf("%d %d", &from, &to);
    G[from].push_back(to);
    G[to].push_back(from);
  }

  for (int i = 0; i < first_N; i++) {

    chain_max_depth[i] = -1; // 鎖の中の最大の深さ
  }

  root = 0;
  makeLA();
  makeTree(root, -1, 0);
  global_index = 0;
  HLD(root, root, -1, 0); // 0スタートでdfs木を作ったので、0がrootとなる

  // dfs木をセグメント木に変換する
  N = 1;
  init(first_N);

  for (int i = 0; i < N; i++) {

    POOL[i] = 0;
    uniformity[i] = 0;
    partial[i] = 0;
  }

  int command, node_id;
  ll add_weight;

  for (int loop = 0; loop < num_query; loop++) {

    scanf("%d", &command);

    if (command == 0) { // 重みの和を計算

      scanf("%d %d", &from, &to);
      printf("%lld\\n", calc(from, to));

    } else {

      scanf("%d %lld", &node_id, &add_weight);
      add(node_id, add_weight);
    }
  }

  return 0;
}