[LeetCode] 2563. Count the Number of Fair Pairs

🔈 문제

Given a 0-indexed integer array nums of size n and two integers lower and upper, return the number of fair pairs.
A pair (i, j) is fair if:
0 <= i < j < n, and lower <= nums[i] + nums[j] <= upper

 

💫 예시

Input: nums = [0,1,7,4,4,5], lower = 3, upper = 6
Output: 6
Explanation: There are 6 fair pairs: (0,3), (0,4), (0,5), (1,3), (1,4), and (1,5).

 

🎱 제한

• 1 <= nums.length <= $10^5$
• nums.length == n
• $-10^9$ <= nums[i] <= $10^9$
• $-10^9$ <= lower <= upper <= $10^9$

 

📚 문제 풀이

 

문제를 요약하자면 nums에서 i, j 를 골라(i != j)  더했을 때 lower 보다 같거나 크고, upper 보다 작거나 같은 짝꿍 수를 구하는 문제이다. 

 

이중 for문 돌려가며 구하면 $O(n^2)$로 시간초과가 날 것이다. 그러면 좀 더 효율적인 알고리즘을 선택해야 하는데 여기서 떠오르는 것은 이분탐색이다. 숫자를 고를 때에는 서로 다른 숫자만 고르면 되는 것이지 다른 추가 제한사항은 없었다. 

 

우선 nums를 정렬하고, nums[i]에 대해서 lower <= nums[i] + nums[j] <= upper를 만족하는 nums[j]를 찾아야 한다. 이를 분리해서 보면, lower <= nums[i] + nums[j], nums[i] + nums[j] <= upper 로 나눠서 생각할 수 있는데, 식을 다시 바꿔보면 결국 lower - nums[i] <= nums[j] 와 nums[j] <= upper - nums[i] 가 된다.

 

lower - nums[i] <= nums[j] 를 만족하는 최소 인덱스 j를 left 라고 하고, nums[j] <= upper - nums[i]를 만족하는 최대 인덱스 j를 right라고 했을 때, right - left + 1이 현재 인덱스 i에 대해서 lower <= nums[i] + nums[j] <= upper 를 만족하는 갯수가 된다. 

 

최소 인덱스를 찾기 위해서 lowerBound를 사용하였고, 최대 인덱스를 찾기 위해서 upperBound를 사용하였다. 각각 인덱스는 현재 i 보다 커야 하기 때문에 범위도 직접 지정해주었다. 

 

left랑 right는 이분탐색을 사용하여 구하면 각각 $O(logN)$ 으로 이 문제는 $O(n * logN)$ 으로 해결할 수 있다. 

 

💻 코드

이번에는 go로 문제를 풀어 보았다.

 

func upperBound(nums []int, start int, end int, target int) int {
	left := start
	right := end
	for left <= right {
		mid := left + (right-left)/2
		if nums[mid] > target {
			right = mid - 1
		} else {
			left = mid + 1
		}
	}
	return right
}

func lowerBound(nums []int, start int, end int, target int) int {
	left := start
	right := end
	for left <= right {
		mid := left + (right-left)/2
		if nums[mid] < target {
			left = mid + 1
		} else {
			right = mid - 1
		}
	}
	return left
}

func countFairPairs(nums []int, lower int, upper int) int64 {
	var res int64 = 0
	sort.Ints(nums)
	n := len(nums)
	for i := 0; i < n; i++ {
		left := lowerBound(nums, i+1, n-1, lower-nums[i])
		right := upperBound(nums, i+1, n-1, upper-nums[i])
		if left <= right {
			res += int64(right - left + 1)
		}
	}
	return res
}

 

🔗 문제링크