线段树的 js 实现

Segment tree 可以用来解决一些区间的问题。比如说有 100 个元素的整数数组，想要求其中索引位置在 a - b 之间所有数字的和，那怎么求呢？一种方式是遍历，从 a 到 b，这样的时间复杂度是 O(n) 级别的；第二种是使用线段树，则可以把时间复杂度优化到 O(logn)。

下面是 js 的线段树实现，可以直接拷贝使用。使用方法请往下面看。

function SegmentTree(aggregate) {
  this._data = [];
  this._original = null;
  this._aggregate = aggregate;
}
/**
 * Creates a segment tree using an array passed as element.
 *
 * @static
 * @public
 * @param {Array} array Array to be indexed.
 * @param {Function} aggregate Function used for generation of
 *  intermediate nodes.
 */
SegmentTree.indexArray = function (array, aggregate) {
  var segmentize = function (original, data, lo, hi, idx) {
    if (lo === hi) {
      data[idx] = original[lo];
      return;
    }

    var mid = Math.floor((lo + hi) / 2);
    var left = 2 * idx + 1;
    var right = 2 * idx + 2;

    segmentize(original, data, lo, mid, left);
    segmentize(original, data, mid + 1, hi, right);

    data[idx] = aggregate(data[left], data[right]);
  };

  var result = new Array(array.length * 4);
  if (array && array.length) {
    segmentize(array, result, 0, array.length - 1, 0);
  }
  var tree = new SegmentTree(aggregate);
  tree._data = result;
  tree._original = array;
  return tree;
};

SegmentTree.prototype._leftChild = function (index) {
  return 2 * index + 1;
};

SegmentTree.prototype._rightChild = function (index) {
  return 2 * index + 2;
};

SegmentTree.prototype.update = function (index, v) {
  this._original[index] = v;
  let update = (treeIndex, l, r, index, v) => {
    if (l === r) {
      this._data[treeIndex] = v;
      return;
    }

    let leftTreeIndex = this._leftChild(treeIndex);
    let rightTreeIndex = this._rightChild(treeIndex);
    let mid = l + (((r - l) / 2) | 0);

    if (index >= mid + 1) {
      update(rightTreeIndex, mid + 1, r, index, v);
    } else {
      // index <= mid
      update(leftTreeIndex, l, mid, index, v);
    }

    this._data[treeIndex] = this._aggregate(
      this._data[leftTreeIndex],
      this._data[rightTreeIndex]
    );
  };

  update(0, 0, this._original.length - 1, index, v);
};

/**
 * Queries the SegmentTree in given range based on the set aggregate.
 *
 * @param {Number} start The start index of the interval.
 * @param {Number} end The end index of the interval.
 */
SegmentTree.prototype.query = function (start, end) {
  if (start > end)
    throw new Error("The start index should be smaller by the end index");

  let findEl = (
    originalArrayStart,
    originalArrayEnd,
    current,
    queryLeft,
    queryRight
  ) => {
    if (queryLeft === originalArrayStart && queryRight === originalArrayEnd)
      return this._data[current];

    let originalArrayMid = Math.floor(
      (originalArrayStart + originalArrayEnd) / 2
    );

    let left = 2 * current + 1;
    let right = 2 * current + 2;

    if (queryLeft >= originalArrayMid + 1) {
      return findEl(
        originalArrayMid + 1,
        originalArrayEnd,
        right,
        queryLeft,
        queryRight
      );
    }

    if (queryRight <= originalArrayMid) {
      return findEl(
        originalArrayStart,
        originalArrayMid,
        left,
        queryLeft,
        queryRight
      );
    }

    let leftResult = findEl(
      originalArrayStart,
      originalArrayMid,
      left,
      queryLeft,
      originalArrayMid
    );
    let rightResult = findEl(
      originalArrayMid + 1,
      originalArrayEnd,
      right,
      originalArrayMid + 1,
      queryRight
    );

    return this._aggregate(leftResult, rightResult);
  };

  return findEl(0, this._original.length - 1, 0, start, end);
};

How to use

let arr = [-2, 0, 3, -5, 2, -1]; // 需要构建线段树的数组
let segmentTree = SegmentTree.indexArray(arr, (a, b) => a + b);

console.log(segmentTree.query(0, 2)); // 1
console.log(segmentTree.query(2, 5)); // -1
console.log(segmentTree.query(0, 5)); // -3
console.log(segmentTree);

segmentTree.update(1, 20);

console.log(segmentTree.query(0, 5)); // 17