Segment Tree implementation (trekhleb#45)

vignesh-m · trekhleb · commit 2e5fd8d583e7 · 2018-06-04T06:32:34.000+03:00
* added segment tree implementation - supports custom operation

* added readme for segment tree
diff --git a/src/data-structures/tree/segment-tree/README.md b/src/data-structures/tree/segment-tree/README.md
@@ -0,0 +1,9 @@
+# Segment Tree
+
+A segment tree is a data structure designed to perform certain array operations efficiently - especially those involving range queries. A common application is the [Range Minimum Query](https://en.wikipedia.org/wiki/Range_minimum_query) (RMQ) problem, where we are given an array of numbers and need to support operations of updating values of the array and finding the minimum of a contiguous subarray. A segment tree implementation for the RMQ problem takes O(n) to initialize, and O(log n) per query or update. The "minimum" operation can be replaced by any array operation (such as sum).
+
+A segment tree is a binary tree with contiguous subarrays as nodes. The root of the tree represents the whole array. The two children of the root represent the first and second halves of the array. Similarly, the children of each node corresponds to the two halves of the array corresponding to the node. If the array has size n, we can prove that the segment tree has size at most 4n. Each node stores the minimum of its corresponding subarray.
+
+In the implementation, we do not explicity store this tree structure, but represent it using a $4n$ sized array. The left child of node i is 2i and the right child is 2i+1. This is a standard way to represent segment trees, and lends itself to an efficient implementation.
+
+We build the tree bottom up, with the value of each node being the minimum of its children's values. This will take time O(n), with one operation for each node. Updates are also done bottom up, with values being recomputed starting from the leaf, and up to the root. The number of operations done is the height of the tree, which is O(log n) To answer queries, each node splits the query into two parts, one subquery for each child. If a query contains the whole subarray of a node, we can use the precomputed value at the node. Using this optimisation, we can prove that only O(log n) minimum operations are done.
diff --git a/src/data-structures/tree/segment-tree/SegmentTree.js b/src/data-structures/tree/segment-tree/SegmentTree.js
@@ -0,0 +1,149 @@
+/**
+ * Segment Tree implementation for Range Query data structure
+ * Tracks a array of numbers. 0 indexed
+ * operation is a binary function (eg sum, min) - needs to be associative
+ * identity is the identity of the operation
+ *    i.e, operation(x, identity) = x (eg 0 for sum, Infinity for min)
+ * Supports methods
+ *  update(index, val) - set value of index
+ *  query(l, r) - finds operation(values in range [l, r]) (both inclusive)
+ *
+ * As is customary, we store the tree implicitly with i being the parent of 2i, 2i+1.
+ */
+
+export default class SegmentTree {
+  /**
+   * array initialises the numbers
+   * @param {number[]} array
+   */
+  constructor(array, operation, identity) {
+    this.n = array.length;
+    this.array = array;
+    this.tree = new Array(4 * this.n);
+
+    this.operation = operation;
+    this.identity = identity;
+
+    // use Range Min Query by default
+    if (this.operation === undefined) {
+      this.operation = Math.min;
+      this.identity = Infinity;
+    }
+
+
+    this.build();
+  }
+
+  /**
+   * Stub for recursive call
+   */
+  build() {
+    this.buildRec(1, 0, this.n - 1);
+  }
+
+  /**
+   * Left child index
+   * @param {number} root
+   */
+  left(root) {
+    return 2 * root;
+  }
+
+  /**
+   * Right child index
+   * @param {number} root
+   */
+  right(root) {
+    return (2 * root) + 1;
+  }
+
+  /**
+   * root is the index in the tree, [l,r] (inclusive) is the current array segment being built
+   * @param {number} root
+   * @param {number} l
+   * @param {number} r
+   */
+  buildRec(root, l, r) {
+    if (l === r) {
+      this.tree[root] = this.array[l];
+    } else {
+      const mid = Math.floor((l + r) / 2);
+      // build left and right nodes
+      this.buildRec(this.left(root), l, mid);
+      this.buildRec(this.right(root), mid + 1, r);
+      this.tree[root] = this.operation(this.tree[this.left(root)], this.tree[this.right(root)]);
+    }
+  }
+
+  /**
+   * Stub for recursive call
+   * @param {number} lindex
+   * @param {number} rindex
+   */
+  query(lindex, rindex) {
+    return this.queryRec(1, lindex, rindex, 0, this.n - 1);
+  }
+
+  /**
+   * [lindex, rindex] is the query region
+   * [l,r] is the current region being processed
+   * Guaranteed that [lindex,rindex] contained in [l,r]
+   * @param {number} root
+   * @param {number} lindex
+   * @param {number} rindex
+   * @param {number} l
+   * @param {number} r
+   */
+  queryRec(root, lindex, rindex, l, r) {
+    // console.log(root, lindex, rindex, l, r);
+    if (lindex > rindex) {
+      // happens when mid+1 > r - no segment
+      return this.identity;
+    }
+    if (l === lindex && r === rindex) {
+      // query region matches current region - use tree value
+      return this.tree[root];
+    }
+    const mid = Math.floor((l + r) / 2);
+    // get left and right results and combine
+    const leftResult = this.queryRec(this.left(root), lindex, Math.min(rindex, mid), l, mid);
+    const rightResult = this.queryRec(
+      this.right(root), Math.max(mid + 1, lindex), rindex,
+      mid + 1, r,
+    );
+    return this.operation(leftResult, rightResult);
+  }
+
+  /**
+   * Set array[index] to value
+   * @param {number} index
+   * @param {number} value
+   */
+  update(index, value) {
+    this.array[index] = value;
+    this.updateRec(1, index, value, 0, this.n - 1);
+  }
+
+  /**
+   * @param {number} root
+   * @param {number} index
+   * @param {number} value
+   * @param {number} l
+   * @param {number} r
+   */
+  updateRec(root, index, value, l, r) {
+    if (l === r) {
+      // we are at tree node containing array[index]
+      this.tree[root] = value;
+    } else {
+      const mid = Math.floor((l + r) / 2);
+      // update whichever child index is in, update this.tree[root]
+      if (index <= mid) {
+        this.updateRec(this.left(root), index, value, l, mid);
+      } else {
+        this.updateRec(this.right(root), index, value, mid + 1, r);
+      }
+      this.tree[root] = this.operation(this.tree[this.left(root)], this.tree[this.right(root)]);
+    }
+  }
+}
diff --git a/src/data-structures/tree/segment-tree/__test__/SegmentTree.test.js b/src/data-structures/tree/segment-tree/__test__/SegmentTree.test.js
@@ -0,0 +1,136 @@
+import SegmentTree from '../SegmentTree';
+
+describe('SegmentTree', () => {
+  it('create RMQ SegmentTree', () => {
+    const array = [1, 2, 5, 3, 4, 6, 2];
+    const segTree = new SegmentTree(array, Math.min, Infinity);
+
+    expect(segTree.array.sort()).toEqual(array.sort());
+    expect(segTree.n).toBe(7);
+  });
+
+  it('check specific tree indices', () => {
+    const array = [1, 2, 5, 3, 4, 6, 2];
+    const segTree = new SegmentTree(array, Math.min, Infinity);
+
+    // 1 - [0,6]
+    // 2 - [0,3] 3 - [4,6]
+    // 4 - [0,1] 5 - [2,3] 6 - [4,5] 7 - [6,6]
+    // 8 - [0,0] 9 - [1,1] 10 - [2,2] 11 - [3,3] 12 - [4,4] 13 - [5,5]
+    expect(segTree.tree.slice(8, 14)).toEqual(array.slice(0, 6));
+    expect(segTree.tree[7]).toBe(array[6]);
+    expect(segTree.tree[1]).toBe(Math.min(...array));
+    expect(segTree.tree[2]).toBe(Math.min(...array.slice(0, 4)));
+    expect(segTree.tree[6]).toBe(Math.min(...array.slice(4, 6)));
+  });
+
+  it('check another tree for n=8', () => {
+    const array = [5, 4, 2, 1, 4, 1, 3, 1];
+    const segTree = new SegmentTree(array, Math.min, Infinity);
+
+    // 1 - [0,7]
+    // 2 - [0,3] 3 - [4,7]
+    // 4 - [0,1] 5 - [2,3] 6 - [4,5] 7 - [6,7]
+    // 8 - [0,0] 9 - [1,1] 10 - [2,2] 11 - [3,3] 12 - [4,4] 13 - [5,5] 14 - [6,6] 15 - [7,7]
+    expect(segTree.tree.slice(8, 16)).toEqual(array.slice(0, 8));
+    expect(segTree.tree[7]).toBe(Math.min(...array.slice(6, 8)));
+    expect(segTree.tree[1]).toBe(Math.min(...array));
+    expect(segTree.tree[2]).toBe(Math.min(...array.slice(0, 4)));
+    expect(segTree.tree[6]).toBe(Math.min(...array.slice(4, 6)));
+  });
+
+  it('check query', () => {
+    const array = [1, 2, 5, 3, 4, 6, 2];
+    const segTree = new SegmentTree(array, Math.min, Infinity);
+
+    const testRanges = [[0, 6], [0, 4], [2, 6], [3, 3], [4, 5], [6, 6], [1, 5], [1, 4]];
+    for (let i = 0; i < testRanges.length; i += 1) {
+      const range = testRanges[i];
+      expect(segTree.query(range[0], range[1]))
+        .toBe(Math.min(...array.slice(range[0], range[1] + 1)));
+    }
+    expect(segTree.query(0, 0)).toBe(1);
+  });
+
+  it('check update using queries', () => {
+    const array = [1, 2, 5, 3, 4, 6, 2];
+    const segTree = new SegmentTree(array, Math.min, Infinity);
+
+    const testRanges = [[0, 6], [0, 4], [2, 6], [3, 3], [4, 5], [6, 6], [1, 5], [1, 4]];
+
+    expect(segTree.array[0]).toBe(1);
+    for (let i = 0; i < testRanges.length; i += 1) {
+      const range = testRanges[i];
+      expect(segTree.query(range[0], range[1]))
+        .toBe(Math.min(...array.slice(range[0], range[1] + 1)));
+    }
+
+    segTree.update(0, 3);
+    array[0] = 3;
+
+    expect(segTree.array[0]).toBe(3);
+    for (let i = 0; i < testRanges.length; i += 1) {
+      const range = testRanges[i];
+      expect(segTree.query(range[0], range[1]))
+        .toBe(Math.min(...array.slice(range[0], range[1] + 1)));
+    }
+
+    segTree.update(2, 2);
+    array[2] = 2;
+
+    expect(segTree.array[2]).toBe(2);
+    for (let i = 0; i < testRanges.length; i += 1) {
+      const range = testRanges[i];
+      expect(segTree.query(range[0], range[1]))
+        .toBe(Math.min(...array.slice(range[0], range[1] + 1)));
+    }
+  });
+
+  it('check range sum query SegmentTree', () => {
+    const array = [1, 2, 5, 3, 4, 6, 2];
+    const sum = (a, b) => a + b;
+    const segTree = new SegmentTree(array, sum, 0);
+
+    const testRanges = [[0, 6], [0, 4], [2, 6], [3, 3], [4, 5], [6, 6], [1, 5], [1, 4]];
+
+    expect(segTree.array[0]).toBe(1);
+    for (let i = 0; i < testRanges.length; i += 1) {
+      const range = testRanges[i];
+      expect(segTree.query(range[0], range[1]))
+        .toBe(array.slice(range[0], range[1] + 1).reduce(sum));
+    }
+
+    segTree.update(0, 3);
+    array[0] = 3;
+
+    expect(segTree.array[0]).toBe(3);
+    for (let i = 0; i < testRanges.length; i += 1) {
+      const range = testRanges[i];
+      expect(segTree.query(range[0], range[1]))
+        .toBe(array.slice(range[0], range[1] + 1).reduce(sum));
+    }
+  });
+
+  it('check default is rmq', () => {
+    const array = [3, 7, 2, 5, 4, 3, 8, 1];
+    const segTree = new SegmentTree(array);
+
+    const testRanges = [[0, 7], [3, 7], [2, 5], [4, 4]];
+
+    for (let i = 0; i < testRanges.length; i += 1) {
+      const range = testRanges[i];
+      expect(segTree.query(range[0], range[1]))
+        .toBe(Math.min(...array.slice(range[0], range[1] + 1)));
+    }
+
+    segTree.update(0, 1);
+    array[0] = 1;
+
+    expect(segTree.array[0]).toBe(1);
+    for (let i = 0; i < testRanges.length; i += 1) {
+      const range = testRanges[i];
+      expect(segTree.query(range[0], range[1]))
+        .toBe(Math.min(...array.slice(range[0], range[1] + 1)));
+    }
+  });
+});
