LC-327.区间和的个数

题目描述

leetcode 困难题

给你一个整数数组 nums 以及两个整数 lower 和 upper 。求数组中，值位于范围 [lower, upper] （包含 lower 和 upper）之内的 区间和的个数 。

区间和 S(i, j) 表示在 nums 中，位置从 i 到 j 的元素之和，包含 i 和 j (i ≤ j)。

示例1：

输入：nums = [-2,5,-1], lower = -2, upper = 2
输出：3
解释：存在三个区间：[0,0]、[2,2] 和 [0,2] ，对应的区间和分别是：-2 、-1 、2 。

提示1：

1 <= nums.length <= 10^5
-2^31 <= nums[i] <= 2^31 - 1
-10^5 <= lower <= upper <= 10^5
题目数据保证答案是一个 32 位 的整数

树状数组

设前缀和数组为 $preSum$ ，且 $0 <= j < i < preSum.length$ 。

原问题可以转化成对于每个 $preSum[i]$ ，找到满足 $lower <= preSum[i] - preSum[j] <= upper$ 的前缀和元素个数。即对于 $preSum[i]$ 找到位于 $[preSum[i] - upper, perSum[i] - lower]$ 的 $preSum[j]$ 个数。

此时转换成了单点更新、区间查询问题，可以使用树状数组或线段树实现。只需要顺序遍历前缀和数组，查询 $[preSum[i] - upper, perSum[i] - lower]$ 区间的元素个数后，再将当前元素插入树状数组或线段树中即可。

但要注意到的是 $nums[i]$ 范围很大且存在负数，所以我们需要将原数组离散化到连续的整数区间内。

class Solution {
    public int countRangeSum(int[] nums, int lower, int upper) {
        int n = nums.length;
        long[] preSum = new long[n + 1];
        for(int i = 1; i < n + 1; i++){
            preSum[i] = preSum[i - 1] + nums[i - 1];
        }
        long[] sort = Arrays.stream(preSum).sorted().toArray();
        int rank = 1;
        TreeMap<Long, Integer> map = new TreeMap<>();
        for (int i = 0; i < n + 1; i++) {
            if (!map.containsKey(sort[i])) {
                map.put(sort[i], rank++);
            }
        }
        FenwickTree fenwickTree = new FenwickTree(rank);
        int ans = 0;
        for (int i = 0; i < n + 1; i++) {
            Map.Entry<Long, Integer> ceil = map.ceilingEntry(preSum[i] - upper);
            Map.Entry<Long, Integer> floor = map.floorEntry(preSum[i] - lower);
            if (ceil != null && floor != null) { 
                ans += (fenwickTree.query(floor.getValue()) - fenwickTree.query(ceil.getValue() - 1));
            }
            int cur = map.get(preSum[i]);
            fenwickTree.update(cur, 1);
        }
        return ans;
    }
}

class FenwickTree{
    int[] arr;
    int size;

    public FenwickTree(int size) {
        arr = new int[size];
        this.size = size;
    }

    public FenwickTree(int[] org) {
        arr = new int[org.length + 1];
        this.size = arr.length;
        for (int i = 1; i < size; i++) {
            update(i, org[i - 1]);
        }
    }

    public void update(int index, int delta) {
        while (index < size) {
            arr[index] += delta;
            index += lowBit(index);
        }
    }

    public int query(int index) {
        int sum = 0;
        while (index > 0) {
            sum += arr[index];
            index -= lowBit(index);
        }
        return sum;
    }

    private int lowBit(int i) {
        return i & -i;
    }
}

线段树（数组实现）

思路类似上面的树状数组。

class Solution {
    public int countRangeSum(int[] nums, int lower, int upper) {
        int n = nums.length;
        long[] preSum = new long[n + 1];
        for(int i = 1; i < n + 1; i++){
            preSum[i] = preSum[i - 1] + nums[i - 1];
        }
        long[] sort = Arrays.stream(preSum).sorted().toArray();
        int rank = 0;
        TreeMap<Long, Integer> map = new TreeMap<>();
        for (int i = 0; i < n + 1; i++) {
            if (!map.containsKey(sort[i])) {
                map.put(sort[i], rank++);
            }
        }
        SegTree segTree = new SegTree(rank);
        int ans = 0;
        for (int i = 0; i < n + 1; i++) {
            Map.Entry<Long, Integer> ceil = map.ceilingEntry(preSum[i] - upper);
            Map.Entry<Long, Integer> floor = map.floorEntry(preSum[i] - lower);
            if (ceil != null && floor != null) { 
                ans += segTree.query(ceil.getValue(), floor.getValue(), 0, rank - 1, 1);
            }
            int cur = map.get(preSum[i]);
            segTree.update(cur, 1, 0, rank - 1, 1);
        }
        return ans;
    }
}

class SegTree {
    int[] arr;

    public SegTree(int n) {
        this.arr = new int[n * 4];
    }

    public void update(int index, int c, int s, int t, int p){
        if (s == t) {
            arr[p] += c;
            return;
        }
        int mid = s + (t - s >> 1);
        if (index <= mid) {
            update(index, c, s, mid, p * 2);
        } else {
            update(index, c, mid + 1, t, p * 2 + 1);
        }
        arr[p] = arr[p * 2] + arr[p * 2 + 1];
    }

    public int query(int l, int r, int s, int t, int p){
        if (l <= s && r >= t) {
            return arr[p];
        }
        int mid = s + (t - s >> 1);
        int sum = 0;
        if (l <= mid) {
            sum += query(l, r, s, mid, p * 2);
        }
        if (r >= mid + 1) {
            sum += query(l, r, mid + 1, t, p * 2 + 1);
        }
        return sum;
    }

}

线段树（对象引用实现）

思路类似上面的树状数组。

class Solution {
    public int countRangeSum(int[] nums, int lower, int upper) {
        int n = nums.length;
        long[] preSum = new long[n + 1];
        for(int i = 1; i < n + 1; i++){
            preSum[i] = preSum[i - 1] + nums[i - 1];
        }
        long[] sort = Arrays.stream(preSum).sorted().toArray();
        int rank = 0;
        TreeMap<Long, Integer> map = new TreeMap<>();
        for (int i = 0; i < n + 1; i++) {
            if (!map.containsKey(sort[i])) {
                map.put(sort[i], rank++);
            }
        }
        SegNode root = build(0, rank - 1, new int[rank]);
        int ans = 0;
        for (int i = 0; i < n + 1; i++) {
            Map.Entry<Long, Integer> ceil = map.ceilingEntry(preSum[i] - upper);
            Map.Entry<Long, Integer> floor = map.floorEntry(preSum[i] - lower);
            if (ceil != null && floor != null) { 
                ans += root.query(ceil.getValue(), floor.getValue());
            }
            int cur = map.get(preSum[i]);
            root.update(cur, 1);
        }
        return ans;
    }

    public SegNode build(int left, int right, int[] arr){
        SegNode node = new SegNode(left, right);
        if(left == right){
            node.add = arr[left];
            return node;
        }
        int mid = left + (right - left >> 1);
        if(left <= mid){
            node.leftChild = build(left, mid, arr);
        }
        if(mid + 1 <= right){
            node.rightChild = build(mid + 1, right, arr);
        }
        return node;
    }

    class SegNode {
        int left;
        int right;
        SegNode leftChild;
        SegNode rightChild;
        int add;

        public SegNode(int left, int right){
            this.left = left;
            this.right = right;
        }

        public void update(int index, int c){
            if (left == right) {
                add += c;
                return;
            }
            int mid = left + (right - left >> 1);
            if (index <= mid) {
                leftChild.update(index, c);
            } else {
                rightChild.update(index, c);
            }
            add = leftChild.add + rightChild.add;
        }

        public int query(int l, int r){
            if (l <= left && r >= right) {
                return add;
            }
            int mid = left + (right - left >> 1);
            int sum = 0;
            if (l <= mid) {
                sum += leftChild.query(l, r);
            }
            if (r >= mid + 1) {
                sum += rightChild.query(l, r);
            }
            return sum;
        }

    }

}

动态开点线段树

另一种类似的做法是我们可以使用动态开点的线段树，此时就无需对原数组进行离散化。

class Solution {
    public int countRangeSum(int[] nums, int lower, int upper) {
        int n = nums.length;
        long[] preSum = new long[n + 1];
        for (int i = 1; i < n + 1; i++) {
            preSum[i] = preSum[i - 1] + nums[i - 1];
        }
        long leftBound = Integer.MAX_VALUE, rightBound = Integer.MIN_VALUE;
        for (long item : preSum) {
            leftBound = Math.min(Math.min(leftBound, item), Math.min(item - lower, item - upper));
            rightBound = Math.max(Math.max(rightBound, item), Math.max(item - lower, item - upper));
        }
        SegNode root = new SegNode(leftBound, rightBound);
        int ans = 0;
        for (int i = 0; i < n + 1; i++) {
            ans += root.query(preSum[i] - upper, preSum[i] - lower);
            root.update(preSum[i], 1);
        }
        return ans;
    }

    class SegNode {
        long left;
        long right;
        SegNode leftChild;
        SegNode rightChild;
        int add;

        public SegNode(long left, long right) {
            this.left = left;
            this.right = right;
        }

        public void update(long index, int c) {
            if (left == right) {
                add += c;
                return;
            }
            long mid = left + (right - left >> 1);
            if (index <= mid) {
                if (leftChild == null) {
                    leftChild = new SegNode(left, mid);
                }
                leftChild.update(index, c);
            } else {
                if (rightChild == null) {
                    rightChild = new SegNode(mid + 1, right);
                }
                rightChild.update(index, c);
            }
            add = leftChild == null ? 0 : leftChild.add;
            add += rightChild == null ? 0 : rightChild.add;
        }

        public int query(long l, long r) {
            if (l <= left && r >= right) {
                return add;
            }
            long mid = left + (right - left >> 1);
            int sum = 0;
            if (l <= mid && leftChild != null) {
                sum += leftChild.query(l, r);
            }
            if (r >= mid + 1 && rightChild != null) {
                sum += rightChild.query(l, r);
            }
            return sum;
        }

    }

}