leetcode-go/poj/3468.a-simple-problem-with-integers/main.cc

#include <cstdio>

using namespace std;

const int N = 1000000;

struct TNode {
  long long val;
  long long inc;
  int beg;
  int end;
  TNode *l;
  TNode *r;
  int mid() { return beg + (end - beg) / 2; }
};

TNode tree[2 * N];
int tid = 0;

TNode *build_tree(TNode *root, int beg, int end) {
  root->beg = beg;
  root->end = end;
  root->val = 0;
  root->inc = 0;
  if (beg != end) {
    root->l = build_tree(&tree[tid++], beg, root->mid());
    root->r = build_tree(&tree[tid++], root->mid() + 1, end);
  }
  return root;
}

void update(TNode *root, int beg, int end, long long inc) {
  if (root->beg == beg && root->end == end) {
    root->inc += inc;
    return;
  }
  root->val += inc * (end - beg + 1);
  if (root->mid() < beg) return update(root->r, beg, end, inc);
  if (end <= root->mid()) return update(root->l, beg, end, inc);
  update(root->l, beg, root->mid(), inc);
  update(root->r, root->mid() + 1, end, inc);
}

long long query(TNode *root, int beg, int end) {
  if (beg == root->beg && end == root->end)
    return root->inc * (end - beg + 1) + root->val;
  if (root->inc != 0) {
    root->val += root->inc * (root->end - root->beg + 1);
    update(root->l, root->beg, root->mid(), root->inc);
    update(root->r, root->mid() + 1, root->end, root->inc);
  }
  root->inc = 0;
  if (root->mid() < beg) return query(root->r, beg, end);
  if (end <= root->mid()) return query(root->l, beg, end);
  return query(root->l, beg, root->mid()) +
         query(root->r, root->mid() + 1, end);
}

int main() {
  int n, q;
  scanf("%d %d", &n, &q);
  TNode *root = &tree[tid++];
  build_tree(root, 1, n);
  for (int i = 1, v; i <= n; i++) {
    scanf("%d", &v);
    update(root, i, i, v);
  }
  char op[2];
  for (int i = 0, a, b, c; i < q; i++) {
    scanf("%s", op);
    if (op[0] == 'C') {
      scanf("%d %d %d", &a, &b, &c);
      update(root, a, b, c);
    } else {
      scanf("%d %d", &a, &b);
      printf("%lld\n", query(root, a, b));
    }
  }
  return 0;
}