Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

SegmentTree #90

Open
wants to merge 1 commit into
base: master
Choose a base branch
from
Open
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
feat(init): add initial development of SegmentTree.
crhallberg committed Feb 8, 2023
commit ed3bc11d5c6fe0e5016ee43e915977fd6306f86b
114 changes: 114 additions & 0 deletions src/data-structures/SegmentTree.ts
Original file line number Diff line number Diff line change
@@ -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;
18 changes: 18 additions & 0 deletions src/data-structures/SegmentTreeNode.ts
Original file line number Diff line number Diff line change
@@ -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;
4 changes: 4 additions & 0 deletions src/index.ts
Original file line number Diff line number Diff line change
@@ -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
19 changes: 19 additions & 0 deletions test/data-structures/SequenceTree.test.ts
Original file line number Diff line number Diff line change
@@ -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);
});
});