区间和的个数

题目链接

区间和的个数

题目描述

注意点

• 求数组中，值位于范围 [lower, upper] （包含 lower 和 upper）之内的 区间和的个数
• -2^31 <= nums[i] <= 2^31 - 1
• -10^5 <= lower <= upper <= 10^5

解答思路

• 参照题解使用归并排序解决本题，思路为：
  • 假设数组分为了两个升序排序的部分，从数组arr1从左往右取一个元素i，在确定i之后，从数组arr2找到一个区间[idxR1, idxR2]满足arr2区间内的所有值减去arr1[i]的值都处于[lower, upper]内，[idxR1, idxR2]有多少个元素说明有多少个满足题意的区间和，移动元素i重复上述过程也就找到了本次归并的区间和的个数
  • 需要注意的是，计算区间和时元素必须是相连的，所以不能直接对原数组调换顺序，而应该先计算每个位置的前缀和，然后对前缀和进行归并排序，这样arr2的某个前缀和减去arr1的某个前缀和实际上是某一段连续元素的区间和
  • 需要注意的是，为了方便考虑从第一个元素开始的连续数组的区间和，需要将前缀和数组arr的长度设置为n + 1，arr[0] = 0

代码

class Solution {
    public int countRangeSum(int[] nums, int lower, int upper) {
        int n = nums.length;
        long sum = 0;
        // arr[0]方便计算统计仅取右侧部分元素区间和
        long[] arr = new long[n + 1];
        for (int i = 0; i < n; i++) {
            sum += nums[i];
            arr[i + 1] = sum;
        }
        return countMergeRangeSum(arr, lower, upper, 0, n);
    }

    public int countMergeRangeSum(long[] arr, int lower, int upper, int left, int right) {
        if (left >= right) {
            return 0;
        }
        int res = 0;
        int mid = left + ((right - left + 1) >> 1);
        // 左归并并统计左侧区间和个数
        res += countMergeRangeSum(arr, lower, upper, left, mid - 1);
        // 右归并并统计右侧区间和个数
        res += countMergeRangeSum(arr, lower, upper, mid, right);
        int idxR1 = mid, idxR2 = mid;
        // 随着i向右移动，arr[i]增大，则右侧区间和减去arr[i]要满足[lower, upper]条件，idxR1和idxR2的区间也必定向右移动
        for (int i = left; i < mid; i++) {
            while (idxR1 <= right && arr[idxR1] - arr[i] < lower) {
                idxR1++;
            }
            while (idxR2 <= right && arr[idxR2] - arr[i] <= upper) {
                idxR2++;
            }
            res += idxR2 - idxR1;
        }
        // 归并排序
        long[] newArr = mergeArr(arr, left, right, mid);
        for (int i = 0; i < newArr.length; i++) {
            arr[i + left] = newArr[i];
        }
        return res;
    }

    public long[] mergeArr(long[] arr, int left, int right, int mid) {
        long[] newArr = new long[right - left + 1];
        int idx = 0, l = left, r = mid;
        while (l <= mid - 1 || r <= right) {
            if (l > mid - 1) {
                newArr[idx++] = arr[r++];
                continue;
            }
            if (r > right) {
                newArr[idx++] = arr[l++];
                continue;
            }
            if (arr[l] < arr[r]) {
                newArr[idx++] = arr[l++];
            } else {
                newArr[idx++] = arr[r++];
            }
        }
        return newArr;
    }
}

关键点

• 归并排序的思想
• 初始化前缀和的数组后，数组内的任意两个元素相减都是原数组任意一段连续数组的区间和
• 注意考虑包含第一个元素的区间和的情况