[백준/파이썬] #3151

푸는 데 성공했었을지는 모르겠으나(테스트 케이스는 맞음) 시간 초과가 걸렸다. 

import sys

input = sys.stdin.readline

N = int(input())
A = list(map(int, input().split()))

count = 0
A.sort()
left, right = 0, 0
for i in range(N-2):
    left = i+1
    right = N-1
    while left < right:
        if A[i] + A[left] + A[right] == 0:
            count += 1
            left += 1
        elif A[i] + A[left] + A[right] > 0:
            right -= 1
        else:
            left += 1

print(count)

 

(나중에 pypy3으로 다시 해보니 틀렸더라...!) 참고한 풀이를 보니, for loop 시도는 좋았는데 left, right, i의 값이 서로 겹치는 경우를 따로 처리해주지 못해서 틀린 것으로 보였다. 

 

풀이를 참고해서 푼 코드는 다음과 같다. 이 풀이 덕분에 bisect라는 이진 탐색 라이브러리를 새로 알았다. 코드를 보니 cpython으로 되어 있었다. 이런 라이브러리를 적절히 잘 사용하는 것도 코테의 노하우인 것 같다.

 

참고로 python3으로 풀면 시간초과가 나니 pypy3으로 풀어야 한다.

import sys
from bisect import bisect_left

input = sys.stdin.readline

N = int(input())
A = list(map(int, input().split()))

count = 0
A.sort()
left, right = 0, 0
for i in range(N-2):
    left = i+1
    right = N-1
    while left < right:
        result = A[i] + A[left] + A[right]
        if result == 0:
            if A[left] == A[right]:
                count += right - left
            else:
                idx = bisect_left(A, A[right])
                count += right - idx + 1
            left += 1
        elif result > 0:
            right -= 1
        else:
            left += 1

print(count)