From ed3bc11d5c6fe0e5016ee43e915977fd6306f86b Mon Sep 17 00:00:00 2001
From: Chris Hallberg <chris.hallberg@hey.com>
Date: Wed, 8 Feb 2023 11:00:56 -0500
Subject: [PATCH] feat(init): add initial development of SegmentTree.

---
 src/data-structures/SegmentTree.ts        | 114 ++++++++++++++++++++++
 src/data-structures/SegmentTreeNode.ts    |  18 ++++
 src/index.ts                              |   4 +
 test/data-structures/SequenceTree.test.ts |  19 ++++
 4 files changed, 155 insertions(+)
 create mode 100644 src/data-structures/SegmentTree.ts
 create mode 100644 src/data-structures/SegmentTreeNode.ts
 create mode 100644 test/data-structures/SequenceTree.test.ts

diff --git a/src/data-structures/SegmentTree.ts b/src/data-structures/SegmentTree.ts
new file mode 100644
index 0000000..80b680e
--- /dev/null
+++ b/src/data-structures/SegmentTree.ts
@@ -0,0 +1,114 @@
+import SegmentTreeNode from './SegmentTreeNode';
+
+export type MergeFunction<T> = (a: T, b: T) => T;
+
+class SegmentTree<T> {
+  root: SegmentTreeNode<T>;
+  mergeFn: MergeFunction<T>;
+
+  /**
+   * @constructor
+   * @param {Array<T>} seq
+   * @param {MergeFunction<T>} mergeFn A function to combine or choose between two elements, such as sum, min, max, etc.
+   */
+  constructor(seq: Array<T>, mergeFn: MergeFunction<T>) {
+    this.mergeFn = mergeFn;
+    this.root = this._buildImpl(seq, 0, seq.length - 1);
+  }
+
+  /**
+   * @param {Array<T>} seq
+   * @param {number} from The starting index for this part of the tree
+   * @param {number} to The ending index for this part of the tree
+   * @return {SegmentTreeNode<T>}
+   */
+  protected _buildImpl(
+    seq: Array<T>,
+    from: number,
+    to: number,
+  ): SegmentTreeNode<T> {
+    if (from === to) {
+      return new SegmentTreeNode(seq[from], from, to);
+    }
+
+    // Like a binary tree, but based on index order
+    // instead of the value sort order.
+    const mid = Math.floor((from + to) / 2);
+
+    const left = this._buildImpl(seq, from, mid);
+    const right = this._buildImpl(seq, mid + 1, to);
+
+    const node = new SegmentTreeNode<T>(
+      this.mergeFn(left.value, right.value),
+      from,
+      to,
+    );
+
+    node.left = left;
+    node.right = right;
+
+    return node;
+  }
+
+  /**
+   * Get the result of the mergeFn for a range of indexes.
+   * @param {number} from Inclusive start index
+   * @param {number} to Exclusive end index
+   * @return {T | null}
+   */
+  query(from: number, to: number): T | null {
+    return this._queryImpl(this.root, from, to);
+  }
+
+  /**
+   * Get the result of the mergeFn for a range of indexes.
+   * @param {SegmentTreeNode<T>} node Current node in the recursive search
+   * @param {number} from Inclusive start index
+   * @param {number} to Exclusive end index
+   * @return {T | null}
+   */
+  protected _queryImpl(
+    node: SegmentTreeNode<T>,
+    from: number,
+    to: number,
+  ): T | null {
+    // Node outside range
+    if (node.to < from || to < node.from) {
+      return null;
+    }
+
+    // Node entirely within range
+    if (from <= node.from && node.to <= to) {
+      return node.value;
+    }
+
+    // Partially with range? Well, let's check on the kids
+    const leftValue =
+      node.left !== null
+        ? this._queryImpl(node.left as SegmentTreeNode<T>, from, to)
+        : null;
+    const rightValue =
+      node.right !== null
+        ? this._queryImpl(node.right as SegmentTreeNode<T>, from, to)
+        : null;
+
+    if (leftValue === null) {
+      return rightValue;
+    }
+
+    if (rightValue === null) {
+      return leftValue;
+    }
+
+    return this.mergeFn(leftValue, rightValue);
+  }
+}
+
+/**
+ * Thanks to LeetCode for the example and code that
+ * helped me finally understand this structure!
+ *
+ * https://leetcode.com/articles/a-recursive-approach-to-segment-trees-range-sum-queries-lazy-propagation/
+ */
+
+export default SegmentTree;
diff --git a/src/data-structures/SegmentTreeNode.ts b/src/data-structures/SegmentTreeNode.ts
new file mode 100644
index 0000000..f73eb6a
--- /dev/null
+++ b/src/data-structures/SegmentTreeNode.ts
@@ -0,0 +1,18 @@
+import BinaryTreeNode from './BinaryTreeNode';
+
+class SegmentTreeNode<T> extends BinaryTreeNode<T> {
+  public from: number;
+  public to: number;
+
+  /**
+   * Initialize a sequence tree node with a value.
+   * @param {*} value The value stored in the binary tree node.
+   */
+  constructor(value: T, from: number, to: number) {
+    super(value);
+    this.from = from;
+    this.to = to;
+  }
+}
+
+export default SegmentTreeNode;
diff --git a/src/index.ts b/src/index.ts
index 071eb0e..70ac9c6 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -32,6 +32,8 @@ import NDArray from './data-structures/NDArray';
 import Node from './data-structures/Node';
 import PriorityQueue from './data-structures/PriorityQueue';
 import Queue from './data-structures/Queue';
+import SegmentTree from './data-structures/SegmentTree';
+import SegmentTreeNode from './data-structures/SegmentTreeNode';
 import Stack from './data-structures/Stack';
 import Trie from './data-structures/Trie';
 
@@ -74,6 +76,8 @@ export {
   Node,
   PriorityQueue,
   Queue,
+  SegmentTree,
+  SegmentTreeNode,
   Stack,
   Trie,
   // Utils
diff --git a/test/data-structures/SequenceTree.test.ts b/test/data-structures/SequenceTree.test.ts
new file mode 100644
index 0000000..4d7e93c
--- /dev/null
+++ b/test/data-structures/SequenceTree.test.ts
@@ -0,0 +1,19 @@
+import SegmentTree from './SegmentTree';
+
+describe('search()', () => {
+  test('returns correct result based on value', () => {
+    // Wikipedia example
+    const maxTree = new SegmentTree<number>(
+      [6, 10, 5, 2, 7, 1, 0, 9],
+      Math.max,
+    );
+    expect(maxTree.query(0, 5)).toBe(10);
+
+    // LeetCode example
+    const sumTree = new SegmentTree<number>(
+      [18, 17, 13, 19, 15, 11, 20, 12, 33, 25],
+      (a: number, b: number): number => a + b,
+    );
+    expect(sumTree.query(2, 8)).toBe(123);
+  });
+});