Array-based Segment Tree (just like Heap) solution in Java


  • 0
    M
    /*
        array-based segment tree approach.
        -- suppose input array length is n, need 2 * 2^(ceiling(lgn)) - 1 slots to store the segment sums.
        -- suppose parent's index is i, then its left child's index is 2 * i + 1, right child's index is 2 * i + 2.
    */
    public class NumArray {
    
        int[] tree = null;  // segment sums
        int inputSize = 0;
    
        public NumArray(int[] nums) {
            if (nums != null && nums.length != 0) {
                inputSize = nums.length;
                tree = new int[4 * inputSize - 1];  // 4 * nums.length - 1 > 2 * 2^(ceiling(lgn)) - 1
                build(nums, 0, 0, inputSize - 1);
            }
        }
        
        private void build(int[] nums, int nodeIndex, int segStart, int segEnd) {
            if (segStart == segEnd) {  // leaf node
                tree[nodeIndex] = nums[segStart];
                return;
            }
            int mid = segStart + (segEnd - segStart) / 2;
            build(nums, nodeIndex * 2 + 1, segStart, mid);  // build left subtree
            build(nums, nodeIndex * 2 + 2, mid + 1, segEnd);  // build right subtree
            tree[nodeIndex] = tree[nodeIndex * 2 + 1] + tree[nodeIndex * 2 + 2];
        }
    
        public void update(int i, int val) {
            update(i, val, 0, 0, inputSize - 1);
        }
        
        private void update(int i, int val, int nodeIndex, int segStart, int segEnd) {
            if (segStart == segEnd) {
                tree[nodeIndex] = val;
                return;
            }
            int mid = segStart + (segEnd - segStart) / 2;
            if (i <= mid) update(i, val, nodeIndex * 2 + 1, segStart, mid);
            else update(i, val, nodeIndex * 2 + 2, mid + 1, segEnd);
            tree[nodeIndex] = tree[nodeIndex * 2 + 1] + tree[nodeIndex * 2 + 2];
        }
    
        public int sumRange(int i, int j) {
            return querySum(i, j, 0, 0, inputSize - 1);
        }
        
        private int querySum(int from, int to, int nodeIndex, int segStart, int segEnd) {
            if (from > to) return 0;
            if (from == segStart && to == segEnd) return tree[nodeIndex];
            int mid = segStart + (segEnd - segStart) / 2;
            return querySum(from, Math.min(mid, to), nodeIndex * 2 + 1, segStart, mid) + 
                    querySum(Math.max(mid + 1, from), to, nodeIndex * 2 + 2, mid + 1, segEnd);
        }
    }

Log in to reply
 

Looks like your connection to LeetCode Discuss was lost, please wait while we try to reconnect.