3151번: 합이 0
Elly는 예상치 못하게 프로그래밍 대회를 준비하는 학생들을 가르칠 위기에 처했다. 대회는 정확히 3명으로 구성된 팀만 참가가 가능하다. 그러나 그녀가 가르칠 학생들에게는 큰 문제가 있었다.
www.acmicpc.net
문제
Elly는 예상치 못하게 프로그래밍 대회를 준비하는 학생들을 가르칠 위기에 처했다. 대회는 정확히 3명으로 구성된 팀만 참가가 가능하다. 그러나 그녀가 가르칠 학생들에게는 큰 문제가 있었다. 코딩 실력이 좋으면 팀워크가 떨어지고, 팀워크가 좋을수록 코딩 실력이 떨어진다. 그리고 출전하고자 하는 대회는 코딩 실력과 팀워크 모두가 중요하다.
Elly는 그녀가 가르칠 수 있는 모든 학생들의 코딩 실력을 알고 있다. 각각의 코딩 실력 Ai는 -10000부터 10000 사이의 정수로 표시되어 있다. 그녀는 팀워크와 코딩 실력이 모두 적절한 팀을 만들기 위해, 세 팀원의 코딩 실력의 합이 0이 되는 팀을 만들고자 한다. 이러한 조건 하에, 그녀가 대회에 출전할 수 있는 팀을 얼마나 많이 만들 수 있는지를 계산하여라.
N명의 학생들의 코딩 실력 Ai가 -10000부터 10000사이의 정수로 주어질 때, 합이 0이 되는 3인조를 만들 수 있는 경우의 수를 구하여라.
입력
입력은 표준 입력으로 주어진다.
첫 번째 줄에 학생의 수 N이 입력된다. 두 번째 줄에 N개의 그녀가 가르칠 학생들의 코딩 실력인 Ai가 주어진다.
출력
표준 출력으로 그녀가 고를 수 있는 팀의 수를 하나의 정수로 출력한다.
제한
- 1 ≤ N ≤ 10000
- -10000 ≤ Ai ≤ 10000
코드
from bisect import bisect_left
n = int(input())
arr = sorted(list(map(int, input().split())))
ans = 0
for left in range(n-2):
center, right = left+1, n-1
while center < right:
if arr[left] + arr[center] + arr[right] == 0: # 팀 만들기 가능
if arr[center] == arr[right]: # center~right 다 같은 수
ans += (right - center)
else:
ans += (right - bisect_left(arr, arr[right]) + 1)
center += 1
elif arr[left] + arr[center] + arr[right] < 0:
center += 1
else:
right -= 1
print(ans)
Python3로는 시간 초과로 통과가 안되고 PyPy3로는 통과.
하나의 학생(left)을 고정해놓고 나머지 두 학생(center, right)을 투포인터로 두고 푼다.
만약 세 학생의 코딩 실력의 합이 0이 된다면 2가지 경우를 고려하여 ans에 값을 더해주면 된다.
- arr[center] == arr[right]인 경우, arr 배열은 정렬되어 있기 때문에 center 인덱스부터 right 인덱스까지 같은 값이라는 이야기이다. 그러므로 right-center를 ans에 더해준다
- arr[center] != arr[right]인 경우, arr배열에서 arr[right]와 같은 값을 갖는 인덱스를 찾아서 right부터 그 인덱스까지 길이를 구하여 ans에 더해준다.
만약 세 학생의 코딩 실력의 합이 0보다 작다면 center를 증가, 0보다 크다면 right를 감소시켜 합이 0에 가까워지도록 한다.
'문제풀이 > 투포인터' 카테고리의 다른 글
[Python/파이썬] 백준 20366번 같이 눈사람 만들래? (1) | 2023.04.20 |
---|---|
[Python/파이썬] 백준 22945번 팀 빌딩 (0) | 2023.04.16 |
[Python/파이썬] 백준 17609번 회문 (0) | 2023.03.26 |
[Python/파이썬] 백준 15961번 회전 초밥 (0) | 2023.02.17 |
[Python/파이썬] 백준 22862번 가장 긴 짝수 연속한 부분 수열 (large) (0) | 2023.02.17 |