백준, 합이 0인 네 정수, 7455, 파이썬

[유형]

이진탐색

 

[문제링크]

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

 

[요약]

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

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

 

[문제풀이]

a+b+c+d = 0이므로, a+b = -(c+d)이다.

따라서, dict을 사용해서 a+b를 key에 추가하고, -(c+d)를 key로 가지는 dict이 존재한다면, value만큼 answer에 더해주면 된다.

하지만, 처음에 이 방법을 사용했더니 시간초과가 발생하였다.

import sys
from collections import defaultdict
def input():
    return sys.stdin.readline().rstrip()

n = int(input())

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

vocab_ab = defaultdict()
for a in A:
    for b in B:
        x = a+b
        if x not in vocab_ab.keys():
            vocab_ab[x] =1
        else:
            vocab_ab[x]+=1

for c in C:
    for d in D:
        x = (c+d)*(-1)
        if x in vocab_ab.keys():
            result += vocab_ab[x]

print(result)

 

solve 함수를 선언한 뒤, 해당 동작을 실행하니 시간초과가 발생하지 않았다.

import sys
from collections import defaultdict
def input():
    return sys.stdin.readline().rstrip()

def solve(n, A, B, C,D):
    answer = 0
    vocab_ab = defaultdict()
    for a in A:
        for b in B:
            x = a+b
            if x not in vocab_ab:
                vocab_ab[x] =1
            else:
                vocab_ab[x]+=1

    for c in C:
        for d in D:
            x = -(c+d)
            if x in vocab_ab:
                answer += vocab_ab[x]
    return answer

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

print(solve(n,A,B,C,D))