Segment Tree

Intro

A segment tree is an augmented binary tree. It is designed to store segments or intervals. It is very useful when we want to make interval queries and updates on a linear data structure. There are two ways to implement segment trees. One is the standard object-oriented way which creates a treenode and builds the tree recursively. The other approach is similar to the heap creation which utilizes an array and the indices to simulate a complete binary tree.

Definition

This is a binary tree node so we need pointers to the leftChild and rightChild. To make this tree inherits the segment tree characteristics, we also need an interval and the data associated with the interval.

private class Interval {
    int start, end;
    public Interval(int start, int end) {
        this.start = start;
        this.end = end;
    }
}

private class TreeNode {
    TreeNode leftChild;
    TreeNode rightChild;
    Interval intv;
    Data data;
    public TreeNode(Interval intv, Data data) {
        this.intv = intv;
        this.data = data;
    }
}

Build (LC.201)

Since this is a binary tree, the implementation follows the divide-and-conquer paradigm.

public TreeNode build(int A[], int start, int end) {
    // create node
    TreeNode root = new TreeNode(new Interval(start, end), 0);

    // base case
    if (start == end) {
        root.data = A[start];
        return root;
    }

    // find mid index
    int mid = start + (end - start) / 2;

    // recurse on A[0...mid]
    TreeNode leftSubtree = build(A, start, mid);

    // recurse on A[mid+1...n-1]
    TreeNode rightSubtree = build(A, mid+1, end);

    // merge
    root.leftChild = leftSubtree;
    root.rightChild = rightSubtree;
    root.data = leftSubtree.data + rightSubtree.data;

    return root;
}

Query (LC.202)

We want to divide on the tree data instead of the query data. It's easy to forget which one to use since they have the same variable names.

public int query(SegmentTreeNode root, int start, int end) {
    // base case
    if (root.start == start && root.end == end) {
        return root.max;
    }

    // divide
    int treeMid = root.start + (root.end - root.start) / 2;
    int leftSubtree = Integer.MIN_VALUE, rightSubtree = Integer.MIN_VALUE;

    if (end <= treeMid) { // search left
        leftSubtree = query(root.left, start, end);
    } else if (start > treeMid) { // search right
        rightSubtree = query(root.right, start, end);
    } else { // search both
        leftSubtree = query(root.left, start, treeMid);
        rightSubtree = query(root.right, treeMid+1, end);
    }

    // merge
    return Math.max(leftSubtree, rightSubtree);
}

Modify (LC.203)

Don't forget to update the parent node after updating the leaf node.

public void modify(SegmentTreeNode root, int index, int value) {
    // base case
    if (root.start == index && root.end == index) {
        root.max = value;
        return;
    }

    // divide
    int treeMid = root.start + (root.end - root.start) / 2;
    if (index <= treeMid) {
        modify(root.left, index, value);
    } else {
        modify(root.right, index, value);
    }

    // merge
    root.max = Math.max(root.left.max, root.right.max);
}

References

PreviousK-dimensional TreeNextInterval Sum

Last updated