327. Count of Range Sum #

Problem #

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.

Problem Summary #

Given an integer array nums, return the number of range sums that lie in [lower, upper], including lower and upper. Range sum S(i, j) represents the sum of the elements in nums from position i to j, including i and j (i ≤ j).

Note:
The most straightforward algorithm has complexity O(n^2). Please optimize your algorithm based on this.

Solution Ideas #

• Given an array, the requirement is to find the sum of any subinterval in this array that lies between [lower,upper].

• This problem can be solved by brute force: use 2 nested loops to traverse all subintervals, compute their sums, and determine whether they lie within [lower,upper]. The time complexity is O(n^2).

• Of course, this problem also has a better solution: use a segment tree or a binary indexed tree to reduce the time complexity to O(n log n). The problem requires lower ≤ sum(i,j) ≤ upper, and sum(i,j) = prefixSum(j) - prefixSum(i-1), so lower + prefixSum(i-1) ≤ prefixSum(j) ≤ upper + prefixSum(i-1). Therefore, using prefix sums transforms the range-sum problem into a query problem on prefix sums in a segment tree. However, what is stored in a parent node of the segment tree is not the sum of its child nodes, but the number of occurrences of the child nodes. The second transformation is coordinate compression, because prefix sums can be very large. For example, prefixSum = [-3,-2,-1,0]; coordinate compression is performed using the prefix-sum indices, so the left and right intervals in the segment tree become 0-3.

  

  Use the prefixSum indices for coordinate compression:

  

• Some small details also need attention. After prefixSum is computed, it needs to be deduplicated, and after deduplication it should be sorted, which is convenient for constructing the valid intervals of the segment tree. If it is not deduplicated, invalid intervals (left > right) or overlapping intervals may appear in the segment tree. In the final step, when inserting prefixSum into the segment tree in reverse order, the non-deduplicated version is used. Inserting prefixSum[j] represents the j in sum(i,j). For example, inserting prefixSum[5] into the segment tree means that the case j = 5 has been added to the current tree. The query operation is essentially doing interval matching. For example, when the current i loop reaches i = 3, and prefixSum[5], prefixSum[4], and prefixSum[3] have been inserted into the segment tree, the query operation is essentially checking whether these 3 inequalities hold: lower ≤ sum(i=3,j=3) ≤ upper, lower ≤ sum(i=3,j=4) ≤ upper, lower ≤ sum(i=3,j=5) ≤ upper. The number of inequalities that hold is returned, which is part of the final result to be obtained.

• For example, nums = [-3,1,2,-2,2,-1], prefixSum = [-3,-2,0,-2,0,-1]. After deduplication and sorting, we get sum = [-3,-2,-1,0]. Use coordinate compression to construct the segment tree. For ease of demonstration, the coordinate-compressed segment tree is not shown in the figure below; instead, a non-compressed segment tree is used for display:

  

  Insert len(prefixSum)-1 = prefixSum[5] = -1 in reverse order:

  

  At this point, the search interval becomes [-3 + prefixSum[5-1], -1 + prefixSum[5-1]] = [-3,-1], that is, checking -3 ≤ sum(5,5) ≤ -1. How many cases satisfy the inequality? Here there is obviously only one case, namely j = 5, and it also satisfies the inequality, so in this step res = 1.

• Insert len(prefixSum)-2 = prefixSum[4] = 0 in reverse order:

  

  At this point, the search interval becomes [-3 + prefixSum[4-1], -1 + prefixSum[4-1]] = [-5,-3], that is, checking -5 ≤ sum(4, 4,5) ≤ -3. How many cases satisfy the inequality? Here there are two cases, namely j = 4 or j = 5, and neither satisfies the inequality, so in this step res = 0.

• Insert len(prefixSum)-3 = prefixSum[3] = -2 in reverse order:

  

  At this point, the search interval becomes [-3 + prefixSum[3-1], -1 + prefixSum[3-1]] = [-3,-1], that is, checking -3 ≤ sum(3, 3,4,5) ≤ -1. How many cases satisfy the inequality? Here there are three cases, namely j = 3, j = 4, or j = 5. The cases that satisfy the inequality are j = 3 and j = 5, i.e. -3 ≤ sum(3, 3) ≤ -1 and -3 ≤ sum(3, 5) ≤ -1. So in this step res = 2.

• Insert len(prefixSum)-4 = prefixSum[2] = 0 in reverse order:

  

  At this point, the search interval becomes [-3 + prefixSum[2-1], -1 + prefixSum[2-1]] = [-5,-3], that is, checking -5 ≤ sum(2, 2,3,4,5) ≤ -3. How many cases satisfy the inequality? Here there are four cases, namely j = 2, j = 3, j = 4, or j = 5, and none satisfies the inequality. So in this step res = 0.

• Insert len(prefixSum)-5 = prefixSum[1] = -2 in reverse order:

  

  At this point, the search interval becomes [-3 + prefixSum[1-1], -1 + prefixSum[1-1]] = [-6,-4], that is, checking -6 ≤ sum(1, 1,2,3,4,5) ≤ -4. How many cases satisfy the inequality? Here there are five cases, namely j = 1, j = 2, j = 3, j = 4, or j = 5, and none satisfies the inequality. So in this step res = 0.

• Insert len(prefixSum)-6 = prefixSum[0] = -3 in reverse order:

  

  At this point, the search interval becomes [-3 + prefixSum[0-1], -1 + prefixSum[0-1]] = [-3,-1]. Note that prefixSum[-1] = 0, that is, checking -3 ≤ sum(0, 0,1,2,3,4,5) ≤ -1. How many cases satisfy the inequality? Here there are six cases, namely j = 0, j = 1, j = 2, j = 3, j = 4, or j = 5. The cases that satisfy the inequality are j = 0, j = 1, j = 3, and j = 5, i.e. -3 ≤ sum(0, 0) ≤ -1, -3 ≤ sum(0, 1) ≤ -1, -3 ≤ sum(0, 3) ≤ -1, and -3 ≤ sum(0, 5) ≤ -1. So in this step res = 4. The final answer is the sum of the results of each step: res = 1 + 0 + 2 + 0 + 0 + 4 = 7.

• This problem can also be solved with a binary indexed tree. Similarly, first transform the problem into an interval Query model. lower ≤ prefixSum(j) - prefixSum(i-1) ≤ upper is equivalent to prefixSum(j) - upper ≤ prefixSum(i-1) ≤ prefixSum(j) - lower, where the value of i is in the interval [0,j-1]. Therefore, the problem can be transformed into: for i taking values in the interval [0,j-1], ask how many values among all values in the array prefixSum[0...j-1] lie within the interval [prefixSum(j) - upper, prefixSum(j) - lower]. In a binary indexed tree, the prefix sum of an interval can be transformed into the difference of the prefix sums of 2 intervals, i.e. Query([i,j]) = Query(j) - Query(i-1). Therefore, for this problem, just enumerate whether each value in the array prefixSum[0...j-1] appears within the specified interval and how many times it appears. The first step is to compute all prefix sums prefixSum(j) and [prefixSum(j) - upper, prefixSum(j) - lower], which will be used later. The second step is sorting and coordinate compression; after compression, the point interval is [1,n]. Finally, query the number of occurrences of the value of array prefixSum(j) within the specified interval. Problems with the same pattern include Problem 315 and Problem 493.

Code #

package leetcode

import (
	"sort"

	"github.com/halfrost/leetcode-go/template"
)

// Solution 1: segment tree, time complexity O(n log n)
func countRangeSum(nums []int, lower int, upper int) int {
	if len(nums) == 0 {
		return 0
	}
	st, prefixSum, sumMap, sumArray, res := template.SegmentCountTree{}, make([]int, len(nums)), make(map[int]int, 0), []int{}, 0
	prefixSum[0], sumMap[nums[0]] = nums[0], nums[0]
	for i := 1; i < len(nums); i++ {
		prefixSum[i] = prefixSum[i-1] + nums[i]
		sumMap[prefixSum[i]] = prefixSum[i]
	}
	// sumArray is the deduplicated version of prefixSum; use sumMap to deduplicate
	for _, v := range sumMap {
		sumArray = append(sumArray, v)
	}
	// Sorting ensures that the intervals in the segment tree have left <= right; if not sorted here, many intervals in the segment tree will be invalid.
	sort.Ints(sumArray)
	// Initialize the segment tree; all values in the nodes are assigned 0, i.e. the count is 0
	st.Init(sumArray, func(i, j int) int {
		return 0
	})
	// Reverse order makes it convenient to find j; sum(i,j) requires j >= i, so traverse in reverse order, with i going from large to small
	for i := len(nums) - 1; i >= 0; i-- {
		// The inserted prefixSum[i] is j
		st.UpdateCount(prefixSum[i])
		if i > 0 {
			res += st.Query(lower+prefixSum[i-1], upper+prefixSum[i-1])
		} else {
			res += st.Query(lower, upper)
		}
	}
	return res
}

// Solution 2: binary indexed tree, time complexity O(n log n)
func countRangeSum1(nums []int, lower int, upper int) int {
	n := len(nums)
	// Compute the prefix sum preSum, and all numbers allNums that will be used in later statistics
	allNums, preSum, res := make([]int, 1, 3*n+1), make([]int, n+1), 0
	for i, v := range nums {
		preSum[i+1] = preSum[i] + v
		allNums = append(allNums, preSum[i+1], preSum[i+1]-lower, preSum[i+1]-upper)
	}
	// Coordinate-compress allNums
	sort.Ints(allNums)
	k := 1
	kth := map[int]int{allNums[0]: k}
	for i := 1; i <= 3*n; i++ {
		if allNums[i] != allNums[i-1] {
			k++
			kth[allNums[i]] = k
		}
	}
	// Traverse preSum and use the binary indexed tree to compute the number of valid intervals corresponding to each prefix sum
	bit := template.BinaryIndexedTree{}
	bit.Init(k)
	bit.Add(kth[0], 1)
	for _, sum := range preSum[1:] {
		left, right := kth[sum-upper], kth[sum-lower]
		res += bit.Query(right) - bit.Query(left-1)
		bit.Add(kth[sum], 1)
	}
	return res
}

// Solution 3: brute force, time complexity O(n^2)
func countRangeSum2(nums []int, lower int, upper int) int {
	res, n := 0, len(nums)
	for i := 0; i < n; i++ {
		tmp := 0
		for j := i; j < n; j++ {
			if i == j {
				tmp = nums[i]
			} else {
				tmp += nums[j]
			}
			if tmp <= upper && tmp >= lower {
				res++
			}
		}
	}
	return res
}