[Python/파이썬] 백준 3151번 합이 0

3151번: 합이 0

 

문제

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에 가까워지도록 한다.