[백준 3151] 합이 0 (Python)

www.acmicpc.net/problem/3151

3151번: 합이 0

 

문제풀이

N, arr : 입력값

left, right : 투 포인터

goal : 특정 1개의 원소를 선택하고, 그 수의 음수값

mx_idx : 정렬 후 arr[right]와 동일한 수를 가진 원소의 최소 인덱스

 

세 용액 문제의 응용이라고 생각한다. 알고나면 어렵지 않은 문제인데, 저 알기까지의 과정이 너무 고통스러웠다 ㅠㅠ.. 특정 1개의 원소를 선택 후, 투 포인터(Two-Pointer)를 통해서 세 원소의 합이 0인 경우를 찾아서 경우의 수를 더해주면 된다. 이때 arr[left]가 arr[right]와 동일한지 아닌지에 따라서 다르게 경우의 수를 더해주어야 한다. => O(N^2) 해결

 

처음에는 이분탐색으로 해결하려 하였는데, 두 원소를 선택하는 것에서 O(N^2), 나머지 하나의 원소를 이분탐색으로 찾는 과정에서 O(logN) 이라 전체 시간복잡도는 O(N^2logN) 이므로 해결 불가능하다. (N 조건이 10000 까지)

(혹시 이분탐색으로 풀이가 가능하다면, 댓글로 알려주시면 감사하겠습니다! ㅠㅠ)

 

1. 정렬 후 0부터 N - 2 까지의 반복문을 통해 한 개의 원소 선택하기

- 세 수의 합이 0이 되려면 나머지 두 수의 합이 선택한 원소(arr[i])의 음수값 이어야 한다.

 

2. 투 포인터를 통해 두 원소의 합(goal)이 선택한 원소의 음수(goal) 인지 탐색하기

- tmp < goal 이면 tmp 를 증가하여야 하므로 left 1 증가

- tmp > goal 이면 tmp 를 감소하여야 하므로 right 1 감소

- tmp == goal 일 경우, arr[left]와 arr[right]가 같은지 아닌지 판단하기

- 같을 경우, 가능한 경우의 수는 right - left  ex ) arr = [-4 2 2 2]

- 다를 경우, 가능한 경우의 수는 arr[right]와 동일한 수의 개수 ex) arr = [-5 0 5 5 5]

- tmp == goal 일 경우에도, 정렬을 하였기 때문에 arr[left]와 같은 수가 존재할 수 있으므로 left를 1 증가

 

3. 1, 2 과정을 반복한다.

 

코드

import sys

if __name__ == '__main__':
    N = int(input())
    arr = list(map(int, sys.stdin.readline().split()))
    arr.sort()
    ans = 0
    # 한 개의 수 선택 후 투 포인터를 통해 나머지 두개의 수 구하기
    for i in range(N - 2):
        left, right = i + 1, N - 1
        goal = -arr[i]
        mx_idx = N
        while left < right:
            tmp = arr[left] + arr[right]
            if tmp < goal:
                left += 1
            elif tmp == goal: # 합이 0 인 경우
                # 정렬이 되어있기 때문에 left 원소와 right 원소가 같으면 둘 사이의 거리가 가능한 경우의 수
                if arr[left] == arr[right]:
                    ans += right - left
                else:
                    # left 원소와 right 원소가 서로 다른 경우
                    if mx_idx > right:
                        mx_idx = right
                        # right 에서 1씩 줄이면서 arr[right]와 몇개 있는지 확인
                        while mx_idx >= 0 and arr[mx_idx - 1] == arr[right]:
                            mx_idx -= 1
                    # arr[right]와 동일한 수 갯수 만큼 더하기
                    ans += right - mx_idx + 1
                left += 1
            else:
                right -= 1
    print(ans)

 

결과

 

문제풀이나 코드에서 이상한 부분 지적은 언제나 환영합니다!