[백준 3151] 합이 0 (Python)
3151번: 합이 0
Elly는 예상치 못하게 프로그래밍 대회를 준비하는 학생들을 가르칠 위기에 처했다. 대회는 정확히 3명으로 구성된 팀만 참가가 가능하다. 그러나 그녀가 가르칠 학생들에게는 큰 문제가 있었다.
www.acmicpc.net
문제풀이
N, arr : 입력값
left, right : 투 포인터
goal : 특정 1개의 원소를 선택하고, 그 수의 음수값
mx_idx : 정렬 후 arr[right]와 동일한 수를 가진 원소의 최소 인덱스
세 용액 문제의 응용이라고 생각한다. 알고나면 어렵지 않은 문제인데, 저 알기까지의 과정이 너무 고통스러웠다 ㅠㅠ.. 특정 1개의 원소를 선택 후, 투 포인터(Two-Pointer)를 통해서 세 원소의 합이 0인 경우를 찾아서 경우의 수를 더해주면 된다. 이때 arr[left]가 arr[right]와 동일한지 아닌지에 따라서 다르게 경우의 수를 더해주어야 한다. => O(N^2) 해결
처음에는 이분탐색으로 해결하려 하였는데, 두 원소를 선택하는 것에서 O(N^2), 나머지 하나의 원소를 이분탐색으로 찾는 과정에서 O(logN) 이라 전체 시간복잡도는 O(N^2logN) 이므로 해결 불가능하다. (N 조건이 10000 까지)
(혹시 이분탐색으로 풀이가 가능하다면, 댓글로 알려주시면 감사하겠습니다! ㅠㅠ)
1. 정렬 후 0부터 N - 2 까지의 반복문을 통해 한 개의 원소 선택하기
- 세 수의 합이 0이 되려면 나머지 두 수의 합이 선택한 원소(arr[i])의 음수값 이어야 한다.
2. 투 포인터를 통해 두 원소의 합(goal)이 선택한 원소의 음수(goal) 인지 탐색하기
- tmp < goal 이면 tmp 를 증가하여야 하므로 left 1 증가
- tmp > goal 이면 tmp 를 감소하여야 하므로 right 1 감소
- tmp == goal 일 경우, arr[left]와 arr[right]가 같은지 아닌지 판단하기
- 같을 경우, 가능한 경우의 수는 right - left ex ) arr = [-4 2 2 2]
- 다를 경우, 가능한 경우의 수는 arr[right]와 동일한 수의 개수 ex) arr = [-5 0 5 5 5]
- tmp == goal 일 경우에도, 정렬을 하였기 때문에 arr[left]와 같은 수가 존재할 수 있으므로 left를 1 증가
3. 1, 2 과정을 반복한다.
코드
import sys
if __name__ == '__main__':
N = int(input())
arr = list(map(int, sys.stdin.readline().split()))
arr.sort()
ans = 0
# 한 개의 수 선택 후 투 포인터를 통해 나머지 두개의 수 구하기
for i in range(N - 2):
left, right = i + 1, N - 1
goal = -arr[i]
mx_idx = N
while left < right:
tmp = arr[left] + arr[right]
if tmp < goal:
left += 1
elif tmp == goal: # 합이 0 인 경우
# 정렬이 되어있기 때문에 left 원소와 right 원소가 같으면 둘 사이의 거리가 가능한 경우의 수
if arr[left] == arr[right]:
ans += right - left
else:
# left 원소와 right 원소가 서로 다른 경우
if mx_idx > right:
mx_idx = right
# right 에서 1씩 줄이면서 arr[right]와 몇개 있는지 확인
while mx_idx >= 0 and arr[mx_idx - 1] == arr[right]:
mx_idx -= 1
# arr[right]와 동일한 수 갯수 만큼 더하기
ans += right - mx_idx + 1
left += 1
else:
right -= 1
print(ans)
결과
문제풀이나 코드에서 이상한 부분 지적은 언제나 환영합니다!