[LeetCode] 2563. Count the Number of Fair Pairs - Python

 

문제 설명

---

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

 

문제 예제

---

Example 1:

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).

 

Example 2:

Input: nums = [1,7,9,2,5], lower = 11, upper = 11
Output: 1
Explanation: There is a single fair pair: (2,3).

 

제한 사항

---

• 1 <= nums.length <= 105
• nums.length == n
• -109 <= nums\[i\] <= 109
• -109 <= lower <= upper <= 109

 

내 풀이

class Solution:
    def countFairPairs(self, nums: List[int], lower: int, upper: int) -> int:
        def dfs(nums, now, visited, pair, answer, lower, upper):
            if len(pair) == 2:
                if lower <= now <= upper:
                    answer += 1
                return answer
            
            for i, num in enumerate(nums):
                if not visited[i]:
                    visited[i] = True
                    pair.append(num)

                    answer = dfs(nums, now+num, visited, pair, answer, lower, upper)

                    visited[i] = False
                    if pair: pair.pop()

            return answer
        
        answer = dfs(nums, 0, defaultdict(bool), [], 0, lower, upper)

        return answer//2

 

처음에는 위와 같이 dfs로 접근해서 모든 pair를 탐색하고 조건에 맞으면 answer에 1을 더해주는 코드를 작성했다.

하지만 이 코드는 큰 input 데이터에 대해서 아래 그림과 같이 시간초과가 나는 것을 확인했다.

 

Time Limit Exceeded

 

따라서, 이 문제는 dfs 방식으로 푸는 문제가 아니란 것을 깨닫고, 이것 저것 적어보기 시작했는데, 이진탐색이 떠올랐다. 

 

Input이 [0,1,7,4,4,5]이고 이걸 정렬시켰을 때 [0, 1, 4, 4, 5, 7]이 되니까 각 요소마다 lower를 만족할 수 있는 다른 요소가 있는 위치와 upper를 만족할 수 있는 다른 요소가 있는 위치를 구해 뺀 값을 더해주면 될 것 같았다.

 

이 위치는 이진탐색 알고리즘으로 쉽게 구할 수 있다. 이를 구현한 코드는 다음과 같다.

def binary_search(nums, n, start, is_left):
    left, right = start, len(nums)

    while left < right:
        mid = (left+right)//2

        if nums[mid] < n and is_left:
            left = mid+1
        elif nums[mid] <= n and not is_left:
            left = mid+1
        else:
            right = mid

    return left

class Solution:
    def countFairPairs(self, nums: List[int], lower: int, upper: int) -> int:
        nums.sort()
        n = len(nums)
        answer = 0

        for i in range(n):
            left = binary_search(nums, lower-nums[i], i+1, True)
            right = binary_search(nums, upper-nums[i], i+1, False)
            answer += right - left
            
        return answer

 

이렇게 문제를 해결할 수 있었다. 다른 사람의 풀이를 보니 python에는 bisect라는 이진탐색을 지원하는 함수가 있다는 것을 처음 알게되었다. 알아두면 유용할 것 같은데 이진탐색이 필요할 때 한번 써봐야겠다.

---

[LeetCode] 2563. Count the Number of Fair Pairs