코드트리 | HashMAP / 원소의 합이 0 {언어 : Python}

문제 링크

https://www.codetree.ai/trails/complete/curated-cards/challenge-the-sum-of-the-elements-is-0/description

Code Tree | Learning to Code with Confidence

내 코드

n = int(input())
A = list(map(int, input().split()))
B = list(map(int, input().split()))
C = list(map(int, input().split()))
D = list(map(int, input().split()))

# Write your code here!
answer = 0
dictAB = dict()
dictCD = dict()

def addKey(dictTmp, left, right):
    for i in range(n):
        elemL = left[i]
        for j in range(n):
            elemR = right[j]
            tmp = elemL + elemR
            if dictTmp.get(tmp, 0):
                dictTmp[tmp] += 1
            else:
                dictTmp[tmp] = 1

addKey(dictAB, A, B)
addKey(dictCD, C, D)
keysAB = list(dictAB.keys())
for kab in keysAB:
    kcd = 0 - kab
    if dictCD.get(kcd, 0):
        answer += dictAB[kab] * dictCD[kcd]

print(answer)

더 효율적인 코드

n = int(input())
A = list(map(int,input().split()))
B = list(map(int,input().split()))
C = list(map(int,input().split()))
D = list(map(int,input().split()))

count = dict()
for i in range(n):
    for j in range(n):
        sum = A[i] + B[j]
        if sum in count:
            count[sum] += 1
        else:
            count[sum] = 1

result = 0
for i in range(n):
    for j in range(n):
        target = -(C[i] + D[j])
        if target in count:
            result += count[target]

print(result)

풀이 방법

1. A와 B에서 합을 키로 하고 빈도를 값으로 하는 딕셔너리를 만든다.
2. C와 D에서 합을 키로 하고 빈도를 값으로 하는 딕셔너리를 만든다.
3. 둘의 키를 합쳤을 때 0이 되는 경우 answer에 각 값을 곱해서 더해준다.
   • 예를 들어 AB의 키가 45고 CD의 키가 -45인 경우
   • answer += dictAB[45] * dictCD[-45]
   • 주의할 점
     • dictAB의 키의 개수는 1 <= 키의 개수 <= 5000 * 5000이 된다.
     • 따라서 이중 반복문으로 dictCD에 있는 키 값을 찾으면 안된다.
     • 즉 dictCD에 0 - kab가 존재하는지만 확인하면 되므로 이중 반복문을 사용하지 않아도 된다.
4. answer를 출력한다.

느낀점

• keys값의 개수가 5000개라고 생각해서 2중 for문을 사용한 결과 시간초과가 발생했다.
• 더 효율적인 코드의 경우 두 번째 CD의 합을 구할 때 바로 0이 되는 경우의 수를 구해준다.
• 따라서 반복을 한 번 더 할 필요가 없기 때문에 효율적인 코드가 된다.