[백준] 3151번 - 합이 0 (파이썬)

https://www.acmicpc.net/problem/3151

 

# 합이 0
from bisect import bisect_left

N = int(input())
arr = list(map(int, input().split()))
arr.sort()

answer = 0
for i in range(len(arr)-2):
    left, right = i+1, N-1
    while left < right:
        result = arr[i]+arr[left]+arr[right]
        if result > 0:
            right -= 1
        else:
            if result == 0:
                if arr[left] == arr[right]:
                    answer += right - left
                else:
                    idx = bisect_left(arr, arr[right])
                    answer += right-idx+1
            left += 1

print(answer)

 

투포인터, 이분탐색, bisect
1. 입력을 받고 정렬을 한다.
2. 3명을 고를 것이므로 0부터 N-3까지 (i)
    - 이분탐색 ( left = i+1, right = N-1)
        - result : 세 수의 합
        - 0보다 크면, right -= 1
        - 0보다 작거나 같으면, left += 1
          (중복을 허용하기 때문에, 0일 경우 answer에 중복되는 수만큼 한 번에 더하기)
3. answer 출력

 

투포인터 안에서 이분탐색을 사용했다.

bisect를 이용해서 answer에 더해줄 때, 중복되는 수의 개수를 한 번에 더해주었다.

배열에서의 연속되는 수의 길이는 bisect를 이용해서 구할 수 있다.

from bisect import bisect_left, bisect_right


def count(a, left_value, right_value):
    right_index = bisect_right(a, right_value)
    left_index = bisect_left(a, left_value)
    return right_index - left_index


a = [1, 2, 3, 3, 3, 3, 4, 4, 8, 9]

print(count(a, 4, 4)) # 2
print(count(a, -1, 3)) # 6