leetcode-go/leetcode/hard/327.count-of-range-sum.go

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

func mergeSort(nums []int, lower, upper, beg, end int) int {
	if end-beg <= 1 {
		return 0
	}
	mid := (beg + end) / 2
	m, n := mid, mid
	cnt := mergeSort(nums, lower, upper, beg, mid) + mergeSort(nums, lower, upper, mid, end)
	for i := beg; i < mid; i++ {
		for ; m < end && nums[m]-nums[i] < lower; m++ {
		}
		for ; n < end && nums[n]-nums[i] <= upper; n++ {
		}
		cnt += n - m
	}
	merge(nums, beg, mid, end)
	return cnt
}

func merge(nums []int, beg, mid, end int) {
	sorted := make([]int, end-beg)
	i, j, k := 0, beg, mid
	for ; j < mid && k < end; i++ {
		if nums[j] < nums[k] {
			sorted[i] = nums[j]
			j++
		} else {
			sorted[i] = nums[k]
			k++
		}
	}
	if j != mid {
		copy(sorted[i:], nums[j:mid])
	}
	if k != end {
		copy(sorted[i:], nums[k:end])
	}
	copy(nums[beg:end], sorted)
}