LeetCode每日一题(327. Count of Range Sum)

Given an integer array nums and two integers lower and upper, return the number of range sums that lie in [lower, upper] inclusive.

Range sum S(i, j) is defined as the sum of the elements in nums between indices i and j inclusive, where i <= j.

Example 1:

Input: nums = [-2,5,-1], lower = -2, upper = 2
Output: 3

Explanation: The three ranges are: [0,0], [2,2], and [0,2] and their respective sums are: -2, -1, 2.

Example 2:

Input: nums = [0], lower = 0, upper = 0
Output: 1

Constraints:

• 1 <= nums.length <= 105
• -231 <= nums[i] <= 231 - 1
• -105 <= lower <= upper <= 105
• The answer is guaranteed to fit in a 32-bit integer.

---

我们的目标是从 nums 里找到符合条件的切片， 我们将 nums 平分成两半， left = nums[…m], right = nums[m…], 假设有切片 nums[i…=j]符合条件, 那 i 和 j 就三种情况， 要么都在 left， 要么都在 right， 要么 i 在 left， j 在 right。我们只需要计算 i 在 left， j 在 right 的情况，然后再递归计算都在左边和都在右边的数量， 三者相加就是当前范围内符合条件的切片数量。还有一点就是，我们在计算 i 在 left，j 在 right 的情况的时候， 我们用到了 prefix sum， 比较特别的是我们在生成 prefix sum 之后对它进行了排序， 之所以这样做是为了可以用二分法进行搜索。

---

impl Solution {
    fn find_left(prefix: &Vec<i64>, target: i64) -> usize {
        let mut l = 0;
        let mut r = prefix.len() - 1;
        while l < r {
            let m = (l + r) / 2;
            if prefix[m] < target {
                l = m + 1;
                continue;
            }
            r = m;
        }
        l
    }

    fn find_right(prefix: &Vec<i64>, target: i64) -> usize {
        let mut l = 0;
        let mut r = prefix.len() - 1;
        while l < r {
            let m = (l + r) / 2;
            if prefix[m] <= target {
                l = m + 1;
                continue;
            }
            r = m;
        }
        l
    }
    fn split(nums: &Vec<i64>, lower: i64, upper: i64, l: usize, r: usize) -> i32 {
        if l == r {
            return if nums[l] >= lower && nums[l] <= upper {
                1
            } else {
                0
            };
        }
        let m = (l + r) / 2;
        let mut prefix_sum: Vec<i64> = nums[m + 1..=r]
            .into_iter()
            .scan(0, |s, v| {
                *s += *v;
                Some(*s)
            })
            .collect();
        prefix_sum.sort();
        prefix_sum.insert(0, i64::MIN);
        prefix_sum.push(i64::MAX);
        let mut sum = 0;
        let mut count = 0;
        for i in (l..=m).rev() {
            sum += nums[i];
            let right = Solution::find_right(&prefix_sum, upper - sum);
            let left = Solution::find_left(&prefix_sum, lower - sum);
            count += (right - left) as i32;
        }
        let left_next = Solution::split(nums, lower, upper, l, m);
        let right_next = Solution::split(nums, lower, upper, m + 1, r);
        count + left_next + right_next
    }
    pub fn count_range_sum(nums: Vec<i32>, lower: i32, upper: i32) -> i32 {
        Solution::split(
            &nums.iter().map(|v| *v as i64).collect(),
            lower as i64,
            upper as i64,
            0,
            nums.len() - 1,
        )
    }
}