백준 3151 합이 0 / python

문제 : 3151

풀이 point

세개의 합을 0으로 만드는 경우의 수를 찾는다. 다만 문제는 N이 10000이기 때문에 O(N^3) 풀이는 불가능하다.

그렇다면, 우선 정렬을 통해 입력 받은 숫자를 오름차순으로 만든다.

두개의 합은 투포인터를 통해 쉽게 O(N^2)으로 구할 수 있다. 나머지 하나의 수는 이진탐색을 통해 O(logn)으로 구할 수 있다. 총 시간복잡도는 O(N^2logn)이라 통과 가능하다.

풀이 방법

두 개의 합은 투포인터, 나머지 하나는 이진탐색으로 쉽게 구현할 수 있다.
다만 숫자들이 중복될 수 있기 때문에 만약 -1, 0, 1, 1, 1일 경우 경우의 수는 3개이다.

(-1, 0, 1), (-1, 0, 1), (-1, 0, 1)

그래서 0이 되도록하는 숫자 1이 몇개인지도 세주어야 한다. 나 같은 경우는 1이 등장하는 맨첫번째 인덱스와 1이 등장하지 않는 첫번째 인덱스를(1의 마지막 인덱스) 찾아서 빼주어 1의 개수를 구해주었다.

참고로 빠른 입력을 필수로 해주어야 한다.(안하면 시간 초과)

풀이 코드

#백준, 3151 합이 0

import sys
input = sys.stdin.readline

N = int(input().rstrip())
coding = list(map(int, input().rstrip().split()))

cnt = 0
coding.sort()

def binarySearch(start, k):
    left, right = start, N-1
    lowerBound = N
    while left <= right:
        mid = (left + right)//2
        if coding[mid] >= k:
            lowerBound = mid
            right = mid - 1
        else:
            left = mid + 1

    left, right = start, N - 1
    upperBound = N
    while left <= right:
        mid = (left + right)//2
        if coding[mid] > k:
            upperBound = mid
            right = mid - 1
        else:
            left = mid + 1

    dupliacte = upperBound-lowerBound
    
    return dupliacte if dupliacte > 0 else 0

for i in range(N-2):
    for j in range(i+1, N-1):
        sums = coding[i] + coding[j]
        cnt += binarySearch(j+1, -sums)

print(cnt)

사족

나는 직접 이진 탐색을 구현하였다. 공부를 하는 입장이라 이렇게 구현한 것이므로 실제 대회에서는 미리 구현된 라이브러리 함수를 사용하자.