segmentTree.cc

segmentTree.cc

#include <iostream>

using namespace std;

struct TreeNode {
  int start;
  int end;
  int sum; // or max/min
  TreeNode *left;
  TreeNode *right;
  TreeNode(int start, int end, int sum, TreeNode *left, TreeNode *right)
      : start(start), end(end), sum(sum), left(left), right(right){};
};

TreeNode *buildTree(int start, int end, int* val) {
  if (start == end) {
    return new TreeNode(start, end, val[start], nullptr, nullptr);
  }
  int mid = (end - start) / 2 + start;
  TreeNode *left = buildTree(start, mid, val);
  TreeNode *right = buildTree(mid + 1, end, val);
  return new TreeNode(start, end, left->sum + right->sum, left, right);
}

void updateTree(TreeNode *node, int index, int val) {
  if (node->start == node->end == index) {
    node->sum = val;
    return;
  }
  int mid = (node->end - node->start) / 2 + node->start;
  if (index < mid) {
    updateTree(node->left, index, val);
  } else {
    updateTree(node->right, index, val);
  }
  node->sum = node->left->sum + node->right->sum;
}

int rangeSum(TreeNode *node, int start, int end) {
  if (node->start == start && node->end == end) {
    return node->sum;
  }
  int mid = (end - start) / 2 + start;
  if (end <= mid) {
    return rangeSum(node->left, start, end);
  } else if (start > mid) {
    return rangeSum(node->right, start, end);
  }
  return rangeSum(node->left, start, mid) + rangeSum(node->right, mid + 1, end);
}