[백준] 3151 - 합이 0 Python

https://www.acmicpc.net/problem/3151

3151번: 합이 0

N 명의 학생들의 코딩실력이 -10000부터 10000사이의 정수로 중복을 허용해 주어질 때, 3명의 합이 0 이되는 경우의 수를 구하는 문제다.

정렬 후 한 수를 고정하고, 투 포인터로 첫 번째 수까지 제외한 구간에서 나머지 두 수의 합이 -(첫번째수) 를 만족하도록 만들면 된다.

 

수의 중복을 허용하기 때문에 두 가지의 케이스를 생각해볼 수 있다.

- 합을 0 으로 만드는 두 수가 여러 개 존재하는 경우

- 합을 0 으로 만드는 두 수가 같은 경우

 

첫 번째 케이스는 다음과 같은 경우다.

-7 -7 1 1 1 1 4 5 6 6

합이 0 이되는 경우는 (-7,1,6) 이 존재하는데 1 과 6 이 여러 개 존재한다.이때는 (1 의 개수 * 6의 개수) 가지 케이스가 존재하므로, start 의 요소가 1을 만족하지 않을 때까지 start 를 증가시켜주고, end 의 요소가 6 을 만족시키지 않을 때까지 감소시켜주면 된다.

j,k = start, end
while A[j] == A[start] and j < end:
    j += 1
while A[k] == A[end] and k > start:
    k -= 1
ans += (j-start)*(end-k)

두 번째 케이스는 다음과 같은 경우다.

-10 5 5 5 5 5 5

합이 0 이되는 경우는 (-10,5,5) 가 존재하는데 두 수가 5로 같다.

이때는 5 + 4 + 3 + 2 + 1 = 15 가지 케이스가 존재한다. 이는 start 를 1씩 증가시켰을 때 end-start 의 개수의 합과 같다.

 if A[start] == A[end]:
    ans += (end-start)
    start += 1

전체 코드는 다음과 같다.

 

from sys import stdin

N = int(stdin.readline())
A = list(map(int,stdin.readline().split()))

ans = 0
A.sort()
# 3명으로 구성된 경우만 출전 가능
for i in range(N-2):
    # 투 포인터 알고리즘으로 합이 -A[i] 인 요소들을 구함
    start = i+1
    end = N-1
    while start < end:
        s = A[start] + A[end]
        if s == -A[i]:
            if A[start] == A[end]:
                ans += (end-start)
                start += 1
            else:
                j,k = start, end
                while A[j] == A[start] and j < end:
                    j += 1
                while A[k] == A[end] and k > start:
                    k -= 1
                ans += (j-start)*(end-k)
                start,end = j,k
        elif s < -A[i]:
            start += 1
        else:
            end -= 1
print(ans)