[Python] BOJ / 7453번 / 합이 0인 네 정수

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이 되는 쌍의 개수를 출력한다.

풀이

네 리스트의 조합을 한번에 다루려 할 경우 for문이 4개가 생겨버려 시간 초과가 발생할 가능성이 매우 높아집니다. 따라서 리스트를 둘둘로 나누어 계산해야하는데, 매커니즘은 다음과 같습니다.

 

리스트 A, B, C, D의 각각의 요소 a, b, c, d가 있을 때, (a + b + c + d)는 (a + b) + (c + d)로 표현할 수 있을 것입니다.

 

따라서 먼저 리스트 A와 B의 요소를 더한 새로운 리스트와, C와 D의 요소를 더한 새로운 리스트를 구한 뒤, 이 두 리스트에 대한 계산을 통해 값을 추출할 수 있습니다.

 

이보다 좀 더 이해하기 쉬운 풀이 방법이 있긴 하나 투 포인터 알고리즘을 활용해보고 싶어 해당 방법을 사용하였습니다.

 

자세한 설명은 코드와 함께 주석을 통해 달아놓도록 하겠습니다.

 

참고 블로그

백준 7453번 - 합이 0인 네 정수(★★★ / X▲ / 2) : Python

---

import sys

# 전처리 부분
input = sys.stdin.readline

# 알고리즘 부분
def solve():
    # 투 포인터 알고리즘을 사용하기 위한 인덱스 설정
    # start 변수는 ab 리스트의 첫 번째 인덱스, end 변수는 cd 리스트의 마지막 인덱스를 가리킴
    start, end = 0, len(cd) - 1
    result = 0

    while start < len(ab) and end >= 0:
        if ab[start] + cd[end] == 0:    # 두 포인터가 가리키는 요소의 합이 0이 될 경우
            n_start, n_end = start + 1, end - 1 # 다음 포인터 인덱스 설정
            
            # ab 또는 cd 배열에 중복되는 값이 있을 경우
            while n_start < len(ab) and ab[n_start] == ab[start]:
                n_start += 1
            while n_end >= 0 and cd[n_end] == cd[end]:
                n_end -= 1

            # 전체 쌍의 수를 구하기 위한 연산
            result += (n_start - start) * (end - n_end)

            # 포인터 인덱스 갱신
            start, end = n_start, n_end
        
        # 두 포인터가 가리키는 요소의 합이 0보다 클 경우 or 작을 경우
        elif ab[start] + cd[end] > 0:
            end -= 1    # 합이 0이 되는 요소를 찾기 위해 end 인덱스를 감소시킴
        else: start += 1    # 합이 0이 되는 요소를 찾기 위해 start 인덱스를 증가시킴
    
    print(result)

# 변수 선언 및 초기화 부분
n = int(input())
arr = [list(map(int, input().split())) for _ in range(n)]
ab, cd = [], []

# a와 b, c와 d에 대한 계산 결과를 따로 저장
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()

solve()