Monday, January 26, 2015

Segment Tree (for finding sum of given range) Implementation : Java

Tutorial : http://www.geeksforgeeks.org/segment-tree-set-1-sum-of-given-range/

If the array has N elements.
Then segment tree array has total = 2*N-1 elements.

Number of internal nodes = N-1;


Since we are using array to represent the relations there may be cases when the tree's leaf elements are not all at the maximum depth. If we don't allocate these indices, we'll be getting incorrect results and possible null pointer exceptions. To maintain the tree structure, we must allocate space for missing children as well.



Total space to be allocated thus becomes

    2*(int)Math.pow(2, height)-1;  (and NOT just 2*N-1)

where

Height of segment tree = (int)Math.ceil(Math.log(N)/Math.log(2));

Source:

import java.util.*;

public class SegmentTree {
    private static int[] st;
    
    //Construction of Segment Tree
    public static int[] constructSegmentTree(int[] arr){
        int height = (int)Math.ceil(Math.log(arr.length)/Math.log(2));
        int size = 2*(int)Math.pow(2, height)-1; 
        st = new int[size];
        constructST(arr, 0, arr.length-1, 0);
        return st;
    }
    
    public static int constructST(int[] arr, int ss, int se, int si){
        if(ss==se){
            st[si] = arr[ss];
            return st[si];
        }
        int mid = ss+(se-ss)/2;
        st[si] = constructST(arr, ss, mid, si*2+1)+
                 constructST(arr, mid+1, se, si*2+2);
        return st[si];
    }
    
    //Finding The sum of given Range
    public static int findRangeSum(int ss, int se, int[] arr){
        int n = arr.length;
        if(ss<0 || se > n-1 || ss>se){
            throw new IllegalArgumentException("Invalid arguments");
        }
        return findSum(0, n-1, ss, se, 0);
    }
    
    public static int findSum(int ss, int se, int qs, int qe, int si){ 
        if(ss>qe || se < qs)return 0;
        if(qs<=ss && qe>=se)return st[si];
        
        int mid = ss+(se-ss)/2;
        return findSum(ss, mid, qs, qe, si*2+1)+
               findSum(mid+1, se, qs, qe, si*2+2);
    }
    
    //Updating a particular index's value
    static void updateValue(int arr[], int i, int newVal){
        if(i<0 || i>arr.length-1)throw new IllegalArgumentException();
        int difference = newVal - arr[i];
        arr[i] = newVal;
        updValue(arr, 0, arr.length-1, i, difference, 0);
    }
    
    static void updValue(int[] arr, int ss, int se, int i, int difference, int si){
        if(i< ss || i>se)return;
        st[si] = st[si]+difference;
        if(ss!=se){
            int mid = ss+(se-ss)/2;
            updValue(arr, ss, mid, i, difference, si*2+1);
            updValue(arr, mid+1, se, i, difference, si*2+2);
        }
    } 
    
    public static void main(String[] args) throws Exception{ 
            int[] arr = {1,3,5,7,9,11};
            System.out.println("Height = "+(int)Math.ceil(Math.log(arr.length)/Math.log(2)));
            constructSegmentTree(arr);
            System.out.println("Values array = "+Arrays.toString(arr));
            System.out.println("Segement Tree array = "+Arrays.toString(st));
            System.out.println("Range sum from index 2 to index 4 = "+findRangeSum(2, 4, arr));
            
            //Update value
            updateValue(arr, 2, 2);
            System.out.println("Updated value at index 2 in the original array to "+2);
            System.out.println("Range sum from index 0 to index 2 = "+findRangeSum(0, 2, arr));
    } 
}

Output:

Height = 3
Values array = [1, 3, 5, 7, 9, 11]
Segement Tree array = [36, 9, 27, 4, 5, 16, 11, 1, 3, 0, 0, 7, 9, 0, 0]
Range sum from index 2 to index 4 = 21
Updated value at index 2 in the original array to 2
Range sum from index 0 to index 2 = 6

No comments:

Post a Comment