[알고리즘] BOJ 7453 합이 0인 네 정수 #Python

[BOJ] 7453 합이 0인 네 정수 바로가기

📍 문제

정수로 이루어진 크기가 같은 배열 A, B, C, D가 있다.

A[a], B[b], C[c], D[d]의 합이 0인 (a, b, c, d) 쌍의 개수를 구하는 프로그램을 작성하시오.

---

📍 입력

첫째 줄에 배열의 크기 n (1 ≤ n ≤ 4000)이 주어진다. 다음 n개 줄에는 A, B, C, D에 포함되는 정수가 공백으로 구분되어져서 주어진다. 배열에 들어있는 정수의 절댓값은 최대 228이다.

---

📍 출력

합이 0이 되는 쌍의 개수를 출력한다.

---

📍 풀이

💡 고찰

• 투포인터(Two Pointer) 알고리즘을 적용하여 문제를 해결하였다.
• 너무 많은 경우의 수가 발생할 때 이진 탐색(Binary Search) 혹은 투 포인터(Two Pointer) 알고리즘을 고려해야 한다는 동료의 코멘트가 도움이 되었다.
• 그러나, 아무리 다양한 방법으로 코드를 재 수정하여도 시간초과 문제는 해결할 수 없었다.
• 수 많은 시간초과를 겪고 난 후 다른 훌륭한 개발자 분들의 코드를 참고하여 문제를 해결할 수 있었다.
• 풀이는 틀렸지만 접근 방식은 완전히 엇나가지 않았다는 점이 다행인 것 같다.

📌 문제 풀이

✏️ 1. 4개의 리스트 2개의 리스트로 병합

AB, CD = [], []
for i in range(N):
    for j in range(N):
        AB.append(arr[i][0] + arr[j][1])
        CD.append(arr[i][2] + arr[j][3])
AB.sort()
CD.sort(reverse = True)

• N * N * N * N 의 모든 경우의 수를 구하기엔 시간이 너무 많이 걸린다.
• 경우의 수를 줄이기 위해 4개의 리스트를 2개의 리스트로 병합한다.
  • 리스트 A, B 원소들로 조합 가능한 모든 조합의 합을 리스트 AB에 저장한다.
  • 리스트 C, D 원소들로 조합 가능한 모든 조합의 합을 리스트 CD에 저장한다.
• 병합한 AB, CD 리스트를 각각 오름차순, 내림차순으로 정렬한다.

✏️ 2. 투 포인터 체크

while p1 < length and p2 < length:
    v1, v2 = AB[p1], CD[p2]
    if v1 + v2 == 0:
        n1, n2 = 0, 0
        while v1 == AB[p1] and p1 < length:
            p1, n1 = p1+1, n1+1
            if p1 == length: break
        while v2 == CD[p2] and p2 < length:
            p2, n2 = p2+1, n2+1
            if p2 == length: break
        answer += n1 * n2
    elif v1 + v2 > 0:
        p2 += 1
    elif v1 + v2 < 0:
        p1 += 1

• 투 포인터 알고리즘을 적용하여 원소의 합이 0인 경우를 찾아 answer 에 누적으로 합한다.

✍ 코드

from sys import stdin

def solution(N, arr):
    answer = 0

    # [1] 4개의 리스트 2개로 병합
    AB, CD = [], []
    for i in range(N):
        for j in range(N):
            AB.append(arr[i][0] + arr[j][1])
            CD.append(arr[i][2] + arr[j][3])
    AB.sort()
    CD.sort(reverse = True)

    p1, p2 = 0, 0
    length = N*N

    # [2] 투포인터 체크
    while p1 < length and p2 < length:
        v1, v2 = AB[p1], CD[p2]
        if v1 + v2 == 0:
            n1, n2 = 0, 0
            while v1 == AB[p1] and p1 < length:
                p1, n1 = p1+1, n1+1
                if p1 == length: break
            while v2 == CD[p2] and p2 < length:
                p2, n2 = p2+1, n2+1
                if p2 == length: break
            answer += n1 * n2
        elif v1 + v2 > 0:
            p2 += 1
        elif v1 + v2 < 0:
            p1 += 1
    return answer

# input
N = int(stdin.readline())
arr = list(list(map(int,stdin.readline().split())) for _ in range(N))

# result
result = solution(N, arr)
print(result)