Java实现 LeetCode 307 区域和检索 - 数组可修改

307. 区域和检索 - 数组可修改

给定一个整数数组 nums，求出数组从索引 i 到 j (i ≤ j) 范围内元素的总和，包含 i, j 两点。

update(i, val) 函数可以通过将下标为 i 的数值更新为 val，从而对数列进行修改。

示例:

Given nums = [1, 3, 5]

sumRange(0, 2) -> 9
update(1, 2)
sumRange(0, 2) -> 8
说明:

数组仅可以在 update 函数下进行修改。
你可以假设 update 函数与 sumRange 函数的调用次数是均匀分布的。
PS：
线段树

import java.util.function.*;

public class NumArray{

    private SegmentTree<Integer> segTree;

    public NumArray(int[] nums) {
        if (nums.length>0) {
            Integer[] data=new Integer[nums.length];
            for (int i=0;i<nums.length;i++) {
                data[i]=nums[i];
            }
            segTree = new SegmentTree<Integer>(data,(a,b)->a+b);
        }
    }
    
    public void update(int i, int val) {
        segTree.update(i,val);
    }
    
    public int sumRange(int i, int j) {
        return segTree.searchRange(i,j);
    }
}

class SegmentTree<E>{
    
    private E[] data;

    private E[] tree;

    private BiFunction<E,E,E> function;

    public SegmentTree(E[] arr,BiFunction<E,E,E> function){
        data = (E[]) new Object[arr.length];
        this.function=function;
        System.arraycopy(arr,0,data,0,arr.length);
        tree = (E[]) new Object[4*arr.length];
        buildSegmentTree(0,0,arr.length-1);
    }

    //根据传入的BiFuction构建线段树
    private void buildSegmentTree(int index,int left,int right){
        if (left==right) {
            tree[index] =data[right];
            return;
        }
        int leftIndex=leftChild(index);
        int rightIndex=rightChild(index);
        int mid=left+(right-left)/2;
        buildSegmentTree(leftIndex,left,mid);
        buildSegmentTree(rightIndex,mid+1,right);
        //区间数据和,根据业务需求来
        tree[index]=function.apply(tree[leftIndex],tree[rightIndex]);
    }

    //范围搜索
    public E searchRange(int left,int right){
        return searchRange(0,0,data.length-1,left,right);
    }

    private E searchRange(int rootIndex,int left,int right,int targetLeft,int targetRight){
        if (targetLeft == left && targetRight == right) {
            return tree[rootIndex];
        }
        int mid=left+(right-left)/2;
        if (targetLeft>mid) {
            return searchRange(rightChild(rootIndex),mid+1,right,targetLeft,targetRight);
        }
        if (targetRight<=mid) {
            return searchRange(leftChild(rootIndex),left,mid,targetLeft,targetRight);
        }
        return function.apply(searchRange(leftChild(rootIndex),left,mid,targetLeft,mid),searchRange(rightChild(rootIndex),mid+1,right,mid+1,targetRight));
    }


    public void update(int index,E e){
        if (index<0 || index>=data.length) {
            throw new IllegalArgumentException("index illegal");
        }
        update(0,0,data.length-1,index,e);
    }

    public void update(int rootIndex,int left,int right,int targetIndex,E e){
        if (left == right) {
            tree[rootIndex]=e;
            return;
        }
        int mid=left+(right-left)/2;
        if (targetIndex<=mid) {
            update(leftChild(rootIndex),left,mid,targetIndex,e);
        }else{
            update(rightChild(rootIndex),mid+1,right,targetIndex,e);
        }
        tree[rootIndex]=function.apply(tree[leftChild(rootIndex)],tree[rightChild(rootIndex)]);
    }

    //左孩子
    private int leftChild(int index){
        return index*2+1;
    }

    //右孩子
    private int rightChild(int index){
        return index*2+2;
    }
}