알고리즘[Python]/백준 알고리즘
[ 3151 ] [ 이진 탐색 ] 합이 0
병훈1234
2021. 12. 1. 12:09
https://www.acmicpc.net/problem/3151
3151번: 합이 0
Elly는 예상치 못하게 프로그래밍 대회를 준비하는 학생들을 가르칠 위기에 처했다. 대회는 정확히 3명으로 구성된 팀만 참가가 가능하다. 그러나 그녀가 가르칠 학생들에게는 큰 문제가 있었다.
www.acmicpc.net
풀이 과정
- 임의의 수를 두개 고른 다음, 두 수의 합의 부호를 반전시켜 준다.
- 반전시켜 준 값이 고른 두개의 수 위치 뒤에 있는지 확인한다.
- 있다면, 값이 여러개일 수 있으므로 이분 탐색을 두번 진행한다.
- bisect_left로 왼쪽 위치를 구해주고, bisect_right로 오른쪽 위치를 구해준 다음 두 값의 차이만큼 정답에 더해준다.
- bisect_left의 경우 값이 존재하는 맨 왼쪽 위치가 리턴되고, bisect_right의 경우 값이 존재하는 맨 오른쪽 위치 + 1이 존재하므로 bisect_right의 결과에서 bisect_left의 결과를 빼주면 동일한 값의 개수가 나온다.
소스 코드
import sys
from bisect import bisect_left, bisect_right
input = lambda: sys.stdin.readline().rstrip()
N = int(input())
numbers = list(map(int, input().split()))
numbers.sort()
cache = set(numbers)
answer = 0
for i in range(N):
for j in range(i+1, N):
target = -(numbers[i] + numbers[j])
if target in cache:
answer += bisect_right(numbers, target, j+1, N) - bisect_left(numbers, target, j+1, N)
print(answer)