LeetCode: 327. Count of Range Sum

LeetCode: 327. Count of Range Sum

Given an integer array nums, 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 (i ≤ j), inclusive.

Note:
A naive algorithm of O(n2) is trivial. You MUST do better than that.

Example:

Input: nums = [-2,5,-1], lower = -2, upper = 2,
Output: 3 
Explanation: The three ranges are : [0,0], [2,2], [0,2] and their respective sums are: -2, -1, 2.

解题思路 ——前缀和 + 分治

1. 计算前缀和
2. 根据前缀和递归分治的求解局部的 count
3. 在分治的同时，对前缀和进行排序。这样方便在 O(N) 的复杂度计算同时有分治的两块的数据时，满足要求的序列

AC 代码

func judge(sub int, lower int, upper int) int {
    if sub <= upper && sub >= lower {
        return 1
    } else {
        return 0
    }
}

// [beg, end) 的 Count
func countRangeSumRef(frontSum []int, lower int, upper int) int {
    if len(frontSum) == 0 {
        return 0
    }
    if len(frontSum) == 1 {
        return judge(frontSum[0], lower, upper) 
    }
    
    // divide and conquer
    count := 0
    mid := len(frontSum) / 2
    count += countRangeSumRef(frontSum[:mid], lower, upper)
    count += countRangeSumRef(frontSum[mid:], lower, upper)
   
    // count
    j, k := mid, mid
    for i := 0; i < mid; i++ {
        for j < len(frontSum) && frontSum[j] - frontSum[i] < lower { j++ }
        for k < len(frontSum) && frontSum[k] - frontSum[i] <= upper { k++ }
        count += (k-j)
    }
    
    // sort
    iter := 0
    arr := make([]int, len(frontSum))
    copy(arr, frontSum)
    left, right := 0, mid
    for left < mid || right < len(frontSum) {
        if left >= mid {
            frontSum[iter] = arr[right]
            right++
        } else if right >= len(frontSum) {
            frontSum[iter] = arr[left]
            left++
        } else if arr[left] <= arr[right] {
            frontSum[iter] = arr[left]
            left++
        } else {
            frontSum[iter] = arr[right]
            right++
        }
        iter++
    }

    return count
}

func countRangeSum(nums []int, lower int, upper int) int {
    // 前缀和
    var frontSum []int
    for i := 0; i < len(nums); i++ {
        if i == 0 {
            frontSum =  append(frontSum, nums[i])
        } else {
            frontSum = append(frontSum, frontSum[i-1] + nums[i])
        }
    }

    return countRangeSumRef(frontSum, lower, upper)
}