[leetcode] 2563. Count the Number of Fair Pairs

안녕하세요. 오늘은 leetcode에서 2563. Count the Number of Fair Pairs 문제를 풀어봤습니다.

• https://leetcode.com/problems/count-the-number-of-fair-pairs/description/

 

1. 문제

이 문제에서는 주어진 수열에서 합이 lower 이상, upper 이하인 두 수가 몇개 있는지 찾아야 합니다.

 

예를 들어서 [1, 2, 3, 4, 5] 수열과 lower=4, upper=8이 주어졌다면

 

[(1, 3), (1, 4), (1, 5), (2, 3), (2, 4), (2, 5), (3, 4), (3, 5)] 로 8개가 답인 것이죠.

 

2. 문제 풀이

2.1. brute force

이 문제도 가장 간단하게는 수열 내의 모든 쌍에 대해서 lower 이상, upper 이하인지 검사하는 방법이 있습니다.

• 첫번째 수와 두번째 수
• 첫번째 수와 세번째 수
• ...
• n-1번째 수와 n번째 수

이 방법은 시간복잡도가 O(n^2) 가지게 됩니다.

 

하지만 수열의 길이가 최대 10^5이기 때문에 O(n^2)의 시간복잡도를 가지는 방법으로는 시간 제한에 걸리게 됩니다...!

 

2.2. binary search

그래서 다음으로 생각한 방법이 binary search였습니다.

 

1. 수열을 오름차순으로 정렬합니다.

2. 첫번재 수부터 다음 과정을 수행합니다.

3. i+1~n번째 수 중에서 i번째 수와 더하면 lower보다 크거나 같은 수 중에서 가장 작은 수를 찾습니다.

4. i+1~n번째 수 중에서 i번째 수와 더하면 upper보다 작거나 같은 수 중에서 가장 큰 수를 찾습니다.

5. 3과 4에서 찾은 수들을 더해나갑니다.

 

괜찮은 방법일줄 알았는데 막상 풀고 보니 실행 속도가 느리더라고요 ㅜㅠ

다음에 실행 속도를 개선시켜봐야겠습니다!

부족한 runtime...!

+ 코드

class Solution:
    def countFairPairs(self, nums: List[int], lower: int, upper: int) -> int:
        result = 0
        sorted_nums = sorted(nums)
        for i in range(len(sorted_nums)):
            left_idx = find_smallest_num_greater_than(
                sorted_nums, i, lower - sorted_nums[i]
            )
            right_idx = find_greatest_num_less_than(
                sorted_nums, i, upper - sorted_nums[i]
            )
            if left_idx == -1 or right_idx == -1 or left_idx > right_idx:
                continue
            result += right_idx - left_idx + 1
        return result


def find_greatest_num_less_than(nums, start, num):
    left_idx = start+1
    right_idx = len(nums) - 1
    mid = -1
    
    found = False
    while left_idx <= right_idx:
        mid = (left_idx + right_idx) // 2
        if nums[mid] > num:
            right_idx = mid - 1
        else:
            found = True
            if mid == len(nums) - 1:
                break
            if nums[mid+1] > num:
                break
            else:
                left_idx = mid + 1
    return mid if found else -1

def find_smallest_num_greater_than(nums, start, num):
    left_idx = start+1
    right_idx = len(nums) - 1
    mid = -1
    
    found = False
    while left_idx <= right_idx:
        mid = (left_idx + right_idx) // 2
        if nums[mid] < num:
            left_idx = mid + 1
        else:
            found = True
            if mid == 0:
                break
            if nums[mid-1] < num:
                break
            else:
                right_idx = mid - 1
    return mid if found else -1