2014年11月17日星期一

[Data Structure] Segment Tree

Segment Tree is useful for range sum and update. Given an array A, we have two types of query, one is give the sum of elements in the range [l ,..., r] and the other is update the A[i] with a incremental or decrease. Using Segment Tree, both of the two operation can be done in O(logn) time.

Here is a more detail post about segment tree:http://www.geeksforgeeks.org/segment-tree-set-1-sum-of-given-range/


import java.util.*;

public class SegmentTree {
 public static class TreeNode {
  int val, l ,r;
  TreeNode left, right;
  public TreeNode(int val, int l, int r) {
   this.val = val;
   this.l = l;
   this.r = r;
  }
 }
 
 public static TreeNode buildTree(int[] A, int l, int r) {
  if(l == r) {
   TreeNode node = new TreeNode(A[l], l, r);
   return node;
  }
  else {
   int mid = l + (r - l) / 2;
   TreeNode left = buildTree(A, l, mid);
   TreeNode right = buildTree(A, mid+1, r);
   TreeNode node = new TreeNode(left.val + right.val, l, r);
   node.left = left;
   node.right = right;
   return node;
  }
 }
 
 public static int getSum(TreeNode node, int l, int r) {
  if(node.l >= l && node.r <= r) {
   return node.val;
  }
  else if(node.r < l || node.l > r) {
   return 0;
  }
  else {
   return getSum(node.left, l, r) + getSum(node.right, l, r);
  }
 }
 
 public static void update(TreeNode node, int val, int ind) {
  if(node == null)
   return;
  if(node.l > ind || node.r < ind) {
   return;
  }
  else {
   node.val += val;
   update(node.left, val, ind);
   update(node.right, val, ind);
  }
 }
 
 public static void main(String[] args) {
  Scanner scan = new Scanner(System.in);
  while(scan.hasNext()) {
   int n = scan.nextInt();
   int[] A = new int[n];
   for(int i = 0; i < n; i++)
    A[i] = scan.nextInt();
   TreeNode root = buildTree(A, 0, n-1);
   update(root, 2, 0);
   update(root, 100, 3);
   update(root, -9, 9);
   int l, r;
   while(scan.hasNext()) {
    l = scan.nextInt();
    r = scan.nextInt();
    System.out.println(getSum(root, l, r));
   }
  }
 }
}

没有评论:

发表评论