Interval Sum

Question (LC.206)

Given an array of integers and a list of intervals, return a list of interval sums.

Example

I: [1,2,7,8,5]
Q1: (0, 4) => 1+2+7+8+5
Q2: (1, 2) => 2+7
Q3: (2, 4) => 7+8+5

Approach

1. Build a prefix sum array

2. Write a query interval helper function

Code

public ArrayList<Long> intervalSum(int[] A, ArrayList<Interval> queries) {
    ArrayList<Long> results = new ArrayList<>();
    if (A == null || A.length == 0) {
        return results;
    }
    long[] preSum = prefixSum(A);
    for (Interval query : queries) {
        results.add(getIntvSum(preSum, query));
    }
    return results;
}

// [0, 1, 3, 10, 18, 23]
private long[] prefixSum(int[] nums) {
    long[] preSum = new long[nums.length + 1];
    preSum[0] = 0;
    long curSum = 0;
    for (int i = 0; i < nums.length; i++) {
        curSum += nums[i];
        preSum[i + 1] = curSum;
    }
    return preSum;
}

private long getIntvSum(long[] preSum, Interval query) {
    return preSum[query.end + 1] - preSum[query.start];
}

Time & Space Complexity

The running sum takes O(n) extra space. But the range query only takes O(1) time to process.

Follow Up (LC.207)

Add another function modify that can update the value of A at a given index.

Example

A = [1,2,7,8,5]
query(0, 2) => 10
modify(0, 4) => [4,2,7,8,5]
query(0, 1) => 6
modify(2, 1) => [4,2,1,8,5]
query(2, 4) => 14

Analysis

The prefix sum approach is not so good of an idea now since updating one value will result in updating all the prefix sum after that index. We want a structure that's easy to make updates.

Behold, Segment Tree

The data variable in the segment tree will hold the sum value. Then, we need build the segment tree off the given array. Query and modify should be easy to implement after constructing the tree.

Segment Tree Code

Definition

private class Interval {
    int ST, ED;
    public Interval(int ST, int ED) {
        this.ST = ST;
        this.ED = ED;
    }
}

private class STNode {
    STNode leftChild, rightChild;
    Interval inv;
    long sum;
    public STNode(Interval inv, long sum) {
        this.inv = inv;
        this.sum = sum;
    }
}

Build

private STNode root;

public Solution(int[] A) {
    // init data structure
    if (A == null || A.length == 0) {
        root = null;
    } else {
        root = build(A, new Interval(0, A.length-1));
    }
}

private STNode build(int[] A, Interval inv) {
    // create new node
    STNode node = new STNode(inv, 0L);

    // base case
    if (inv.ST == inv.ED) {
        node.sum = A[inv.ST];
        return node;
    }

    // divide
    int MI = inv.ST + (inv.ED - inv.ST) / 2;
    Interval LInv = new Interval(inv.ST, MI);
    Interval RInv = new Interval(MI+1, inv.ED);

    STNode leftSubtree = build(A, LInv);
    STNode rightSubtree = build(A, RInv);

    // merge
    node.leftChild = leftSubtree;
    node.rightChild = rightSubtree;
    node.sum = leftSubtree.sum + rightSubtree.sum;

    return node;
}

Range Query

/**
 * Return the range sum from index start to index end
 */
public long query(int start, int end) {
    if (root == null || end < start) {
        return -1L;
    }
    return rangeQuery(root, new Interval(start, end));
}

private long rangeQuery(STNode node, Interval query) {
    // base case
    if (node.inv.ST == query.ST && node.inv.ED == query.ED) {
        return node.sum;
    }

    // divide
    int MI = node.inv.ST + (node.inv.ED - node.inv.ST) / 2;
    long leftSubtree = 0L;
    long rightSubtree = 0L;

    if (query.ED <= MI) { // search left
        leftSubtree = rangeQuery(node.leftChild, query);
    } else if (query.ST > MI) { // search right
        rightSubtree = rangeQuery(node.rightChild, query);
    } else { // search both
        Interval leftQuery = new Interval(query.ST, MI);
        Interval rightQuery = new Interval(MI+1, query.ED);
        leftSubtree = rangeQuery(node.leftChild, leftQuery);
        rightSubtree = rangeQuery(node.rightChild, rightQuery);
    }

    // merge
    return leftSubtree + rightSubtree;
}

Update Query

/**
 * Update A[index] to the given value
 */
public void modify(int index, int value) {
    if (root == null || index < 0) {
        return;
    }
    update(root, index, value);
}

private void update(STNode node, int updateIndex, int updateValue) {
    // base case
    if (node.inv.ST == updateIndex && node.inv.ED == updateIndex) {
        node.sum = updateValue;
        return;
    }

    // divide (search left or right)
    int MI = node.inv.ST + (node.inv.ED - node.inv.ST) / 2;

    if (updateIndex <= MI) { // search left
        update(node.leftChild, updateIndex, updateValue);
    } else { // search right
        update(node.rightChild, updateIndex, updateValue);
    }

    // merge (update parent)
    node.sum = node.leftChild.sum + node.rightChild.sum;
}

Time & Space Complexity

The segment tree will still take O(n) extra space. But the range queries and update queries only take O(logn) time.