线段树从零开始

  1 package com.lzp.util;
  2 
  3 import java.util.Arrays;
  4 
  5 /**
  6  * @Author LZP
  7  * @Date 2021/5/5 10:04
  8  * @Version 1.0
  9  *
 10  * 线段树
 11  */
 12 public class SegmentTree {
 13 
 14     /**
 15      * 构造线段树
 16      * @param arr   存储1~n的连续正整数的数组
 17      * @param tree  用于存储构成的线段树的数组（利用满二叉树的形式顺序存储）
 18      * @param root  根节点编号（编号从1开始）
 19      * @param start 左区间（同上）
 20      * @param end   右区间（同上）
 21      */
 22     public static void buildSegTree(int[] arr, int[] tree, int root, int start, int end) {
 23         if (start >= end) {
 24             tree[root - 1] = arr[start - 1];
 25             return;
 26         }
 27 
 28         int mid = (start + end) / 2;
 29         // 左孩子
 30         int left_root  = root * 2;
 31         // 右孩子
 32         int right_root = root * 2 + 1;
 33         // 递归
 34         buildSegTree(arr, tree, left_root, start, mid);
 35         buildSegTree(arr, tree, right_root, mid + 1, end);
 36         // 更新
 37         tree[root - 1] = tree[left_root - 1] + tree[right_root - 1];
 38     }
 39 
 40     /**
 41      * 修改线段树
 42      * 时间复杂度：O(log(n))
 43      * @param arr   存储1~n的连续正整数的数组
 44      * @param tree  用于存储构成的线段树的数组（利用满二叉树的形式顺序存储）
 45      * @param root  根节点编号（编号从1开始）
 46      * @param start 左区间（同上）
 47      * @param end   右区间（同上）
 48      * @param pos   要修改的位置
 49      * @param val   要修改后的值
 50      */
 51     public static void updateSegTree(int[] arr, int[] tree, int root, int start, int end, int pos, int val) {
 52         if (start >= end) {
 53             if (start == pos) {
 54                 arr[start - 1] = val;
 55                 tree[root - 1] = val;
 56             }
 57             return;
 58         }
 59 
 60         int mid = (start + end) / 2;
 61         int left_root  = root * 2;
 62         int right_root = root * 2 + 1;
 63         updateSegTree(arr, tree, left_root, start, mid, pos, val);
 64         updateSegTree(arr, tree, right_root, mid + 1, end, pos, val);
 65         // 更新
 66         tree[root - 1] = tree[left_root - 1] + tree[right_root - 1];
 67     }
 68 
 69     /**
 70      * 查询线段树中区间[L, R]
 71      * 时间复杂度：O(log(n))
 72      * @param arr   存储1~n的连续正整数的数组
 73      * @param tree  用于存储构成的线段树的数组（利用满二叉树的形式顺序存储）
 74      * @param root  根节点编号（编号从1开始）
 75      * @param start 左区间（同上）
 76      * @param end   右区间（同上）
 77      * @param L     区间左边界
 78      * @param R     区间右边界
 79      */
 80     public static int querySegTree(int[] arr, int[] tree, int root, int start, int end, int L, int R) {
 81         // 处理三种特殊情况
 82         if (R < start || L > end) {
 83             // 不相交
 84             return 0;
 85         }
 86         if (L <= start && R >= end) {
 87             // 包含
 88             return tree[root - 1];
 89         }
 90         if (start >= end) {
 91             // 只有一个元素，直接返回
 92             return arr[start - 1];
 93         }
 94 
 95         int mid = (start + end) / 2;
 96         int left_root  = root * 2;
 97         int right_root = root * 2 + 1;
 98         int leftSum  = querySegTree(arr, tree, left_root, start, mid, L, R);
 99         int rightSum = querySegTree(arr, tree, right_root, mid + 1, end, L, R);
100         return leftSum + rightSum;
101     }
102 
103     public static void main(String[] args) {
104         int n = 8;
105         int[] arr = new int[n];
106         for (int i = 0; i < n; i++) {
107             arr[i] = i + 1;
108         }
109         // 线段树
110         // 线段树先将区间[1..n]分成不超过4*n个子区间，对于每个子区间，记录一段连续数字的和。
111         int[] tree = new int[4 * n + 10];
112         buildSegTree(arr, tree, 1, 1, n);
113         System.out.println("建树之后：" + Arrays.toString(tree));
114 
115         updateSegTree(arr, tree, 1, 1, n, 4, 5);
116         System.out.println("修改之后：" + Arrays.toString(tree));
117 
118         // 原本2~7之和为27，但是由于前面已经修改过（把4改成5），所以区间[2, 7]之和为28
119         int left = 2;
120         int right = 7;
121         int result = querySegTree(arr, tree, 1, 1, n, left, right);
122         System.out.printf("区间[%d, %d] = %d\n", left, right, result);
123     }
124 }