문제

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

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

www.acmicpc.net/problem/7453

  • 입력

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

  • 출력

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

최종 코드

앞서 다른 문제를 이분탐색으로 문제를 풀었기때문에 이 문제도 같은 방식으로 접근했다. 다만 이번엔 사전 대신 리스트를 써봤는데 확실히 사전이 편하긴 했다.

근데 python3로는 통과가 안 돼서 다른 접근법이 있는줄 알았는데 알고보니 python3로는 통과한 사람이 없었다.

import sys
from itertools import combinations

n = int(sys.stdin.readline())
A, B, C, D = [], [], [], []
for _ in range(n):
    a, b, c, d = map(int, sys.stdin.readline().split())
    A.append(a)
    B.append(b)
    C.append(c)
    D.append(d)

# A+B / C+D 로 나눔
left = []
right = []
for i in range(n):
    for j in range(n):
        left.append(A[i] + B[j])
        right.append(C[i] + D[j])

# 리스트 정렬
left = sorted(left)
right = sorted(right)
# left는 처음부터 right는 마지막부터
l, r = 0, len(right) - 1
res = 0
while l < len(left) and r >= 0:
    tmp = left[l] + right[r]
    # tmp가 0인 경우 while문을 통해 두 경우의 수 a, b를 각각 구하여 곱해준다
    if tmp == 0:
        l += 1
        r -= 1
        a = b = 1
        while l < len(left) and left[l] == left[l-1]:
            l += 1
            a += 1
        while r >= 0 and right[r] == right[r+1]:
            r -= 1
            b += 1
        res += a*b

    elif tmp < 0:
        l += 1
    else:
        r -= 1

print(res)

다른 방법으로는 Counter 모듈을 사용하는 것이다. 이렇게 구현하니까 훠얼씬 간단하긴 하지만 속도 측면에서는 더 느렸다. 필요할 때에만 시간을 고려해서 사용해야할 것 같다.

import sys
from collections import Counter

n = int(sys.stdin.readline())
A, B, C, D = [], [], [], []
for _ in range(n):
    a, b, c, d = map(int, sys.stdin.readline().split())
    A.append(a)
    B.append(b)
    C.append(c)
    D.append(d)

left = []
right = []
for i in range(n):
    for j in range(n):
        left.append(A[i] + B[j])
        right.append(C[i] + D[j])

counter = Counter(right)
res = 0
for x in left:
    res += counter[-x]

print(res)

참고 사이트