Segment Tree Query II 247

Question

For an array, we can build a SegmentTree for it, each node stores an extra attribute count to denote the number of elements in the the array which value is between interval start and end. (The array may not fully filled by elements)

Design a query method with three parameters root, start and end, find the number of elements in the in array's interval [start, end] by the given root of value SegmentTree.

Example

For array [0, 2, 3], the corresponding value Segment Tree is:

                 [0, 3, count=3]
                 /             \
      [0,1,count=1]             [2,3,count=2]
      /         \               /            \

[0,0,count=1] [1,1,count=0] [2,2,count=1], [3,3,count=1]

query(1, 1), return 0

query(1, 2), return 1

query(2, 3), return 2

query(0, 2), return 2

Solution

和I一样，分情况讨论，递归解决：

1. 查询区间跨越mid，则要分别查询左边区间和右边区间，再将两边结果相加

2. 查询区间只在一边，则只要查询那一边即可

代码如下：

/**
 * Definition of SegmentTreeNode:
 * public class SegmentTreeNode {
 *     public int start, end, count;
 *     public SegmentTreeNode left, right;
 *     public SegmentTreeNode(int start, int end, int count) {
 *         this.start = start;
 *         this.end = end;
 *         this.count = count;
 *         this.left = this.right = null;
 *     }
 * }
 */
public class Solution {
    /**
     *@param root, start, end: The root of segment tree and 
     *                         an segment / interval
     *@return: The count number in the interval [start, end]
     */
    public int query(SegmentTreeNode root, int start, int end) {
        // write your code here
        if(root == null || start > end || start > root.end || end < root.start){
            return 0;
        }

        if((start == root.start && end == root.end) || (root.start == root.end)){
            return root.count;
        }

        int mid = (root.start + root.end) / 2;
        if(start <= mid && end >= mid + 1){
            int left = query(root.left, start, mid);
            int right = query(root.right, mid + 1, end);
            return left + right;
        }else if(end <= mid){
            return query(root.left, start, end);
        }else if(start >= mid + 1){
            return query(root.right, start, end);
        }

        return 0;
    }
}