线段树模板

浅浅的花香味﹌ 2022-10-05 01:53 139阅读 0赞

public class Main { 
    
        public static void main(String[] args) { 
            int[] arr = { 1, 3, 5, 7, 9, 11};
            int[] tree = new int[arr.length * 3];
            buildTree(arr, tree, 0, 0, arr.length - 1);
            updateTree(arr, tree, 0, 4, 6, 0, arr.length - 1);
            System.out.println(query(tree, 0, 0, arr.length - 1, 2, 5));
        }
    
        public static int query(int[] tree, int currentIndex, int start, int end, int rangeLeft, int rangeRight) { 
            if (rangeLeft > end || rangeRight < start) { 
                return 0;
            } else if (start >= rangeLeft && end <= rangeRight) { 
                return tree[currentIndex];
            }
            int midIndex = (start + end) >> 1;
            int leftResult = query(tree, (currentIndex << 1) + 1, start, midIndex, rangeLeft, rangeRight);
            int rightResult = query(tree, (currentIndex << 1) + 2, midIndex + 1, end, rangeLeft, rangeRight);
            return leftResult + rightResult;
        }
    
        public static void updateTree(int[] arr, int[] tree, int currentIndex, int updateIndex, int newVal, int start, int end) { 
            if (updateIndex == start && updateIndex == end) { 
                arr[updateIndex] = newVal;
                tree[currentIndex] = newVal;
                return;
            } else if (updateIndex < start || updateIndex > end) { 
                return;
            }
            int midIndex = (start + end) >> 1;
            int leftNodeIndex = (currentIndex << 1) + 1;
            int rightNodeIndex = (currentIndex << 1) + 2;
            updateTree(arr, tree, leftNodeIndex, updateIndex, newVal, start, midIndex);
            updateTree(arr, tree, rightNodeIndex, updateIndex, newVal, midIndex + 1, end);
            tree[currentIndex] = tree[leftNodeIndex] + tree[rightNodeIndex];
        }
    
        public static void buildTree(int[] arr, int[] tree, int currentIndex, int start, int end) { 
            if (start == end) { 
                tree[currentIndex] = arr[start];
                return;
            }
            int mid = (start + end) >> 1;
            int leftNodeIndex = (currentIndex << 1) + 1;
            int rightNodeIndex = (currentIndex << 1) + 2;
            buildTree(arr, tree, leftNodeIndex, start, mid);
            buildTree(arr, tree, rightNodeIndex, mid + 1, end);
            tree[currentIndex] = tree[leftNodeIndex] + tree[rightNodeIndex];
        }
    
    }