327.区间和的个数

68 阅读1分钟

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

区间和 S(i, j) 表示在 nums 中,位置从 i 到 j 的元素之和,包含 i 和 j (i ≤ j)。
算法:
方法一:归并算法
这里用sums不直接用nums,是为了利用sums[l] - sums[i] < lower和sums[r] - sums[i] <= upper特性。 如果是sums[l] + sums[i] < lower无法确定是移动l还是i

func countRangeSum(nums []int, lower int, upper int) int {
	sums := make([]int, len(nums))
	sum := 0
	for i := range nums {
		sum = sum + nums[i]
		sums[i] = sum
	}
	return mergeSort(sums, 0, len(nums) - 1, lower, upper)
}

func mergeSort(nums []int, left, right, lower, upper int) int {
	if left > right {
		return 0
	}
	if left == right {
		if lower  <= nums[left] && nums[left] <= upper {
			return 1
		} 
		return 0
	}
	mid := (left + right) / 2
	lcount := mergeSort(nums, left, mid, lower, upper)
	rcount := mergeSort(nums, mid + 1, right, lower, upper)
	count := 0 
	// 统计count
	l, r := mid + 1, mid + 1
	for i := left; i <= mid; i ++  {
		for l <= right && nums[l] - nums[i] < lower {
			l ++
		}
		for r <= right && nums[r] - nums[i] <= upper {
			r ++
		}
		count = count + r - l 
	}
	// 排序nums
	tmp := make([]int, 0)
	i, j := left, mid + 1
	for i <= mid && j <= right {
		if nums[i] < nums[j] {
			tmp = append(tmp, nums[i])
			i ++
		} else {
			tmp = append(tmp, nums[j])
			j ++
		}
	}
	for i <= mid {
		tmp = append(tmp, nums[i])
		i ++
	}
	for j <= right  {
		tmp = append(tmp, nums[j])
		j ++
	}
	for i := range tmp {
		nums[left + i] = tmp[i]
	}
	return count + lcount + rcount
}