코드트리 | HashMap / 세 수의 합 {언어 : Python}

문제 링크

https://www.codetree.ai/trails/complete/curated-cards/challenge-sum-of-three-num/description

Code Tree | Learning to Code with Confidence

내 코드

n, k = map(int, input().split())
arr = list(map(int, input().split()))

# Write your code here!
dictNum = dict()
for num in arr:
    if dictNum.get(num, 0):
        dictNum[num] += 1
    else:
        dictNum[num] = 1

answer = 0
keysNum = list(dictNum.keys())
visited = dict()
def makeKeys(a, b, c):
    return [str(a)+str(b)+str(c), 
            str(a)+str(c)+str(b), 
            str(b)+str(a)+str(c),
            str(b)+str(c)+str(a),
            str(c)+str(a)+str(b),
            str(c)+str(b)+str(a)]
lenN = len(keysNum)
for i in range(lenN):
    for j in range(lenN):
        a, b = keysNum[i], keysNum[j]
        c = k - (a + b)
        cntA = dictNum[a]
        cntB = dictNum[b]
        tmp = 0
        if dictNum.get(c, 0):
            cntC = dictNum[c]
            if a == b == c:
                tmp += (cntA * (cntA - 1) * (cntA - 2)) // 6
            elif a == b:
                tmp += ((cntA * (cntA - 1)) // 2) * cntC
            elif a == c:
                tmp += ((cntA * (cntA - 1)) // 2) * cntB
            elif b == c:
                tmp += ((cntB * (cntB - 1)) // 2) * cntA
            else:
                tmp += cntA * cntB * cntC
            vk = makeKeys(a, b, c)
            if not visited.get(vk[0], 0):
                for nk in vk:
                    visited[nk] = 1
                answer += tmp
print(answer)

 

수정 코드

n, k = tuple(map(int, input().split()))
arr = list(map(int, input().split()))

count = dict()

for elem in arr:
    if elem in count:
        count[elem] += 1
    else:
        count[elem] = 1

ans = 0
for i in range(n):
    count[arr[i]] -= 1

    for j in range(i):
        diff = k - arr[i] - arr[j]

        if diff in count:
            ans += count[diff]

print(ans)

풀이 방법

 

1. 빈도 저장: 배열 내 각 숫자의 빈도를 count 딕셔너리에 저장.
2. 조합 탐색:
   • 첫 번째 숫자(arr[i])를 선택 후, 중복 방지를 위해 빈도 감소.
   • 두 번째 숫자(arr[j])를 선택 후, k - arr[i] - arr[j]로 세 번째 숫자(diff) 계산.
   • count에서 diff 존재 여부를 확인하고, 존재하면 빈도만큼 결과 증가.
3. 결과 출력: 가능한 모든 조합의 경우의 수를 ans에 저장하고 출력.

 

느낀점

• 일단 풀 수는 있었지만 빈도를 계산하는 방식이 비효율적이다.
• 따라서 다른 사람의 코드를 확인했고 더 간결하고 시간복잡도가 낮은 코드를 발견했다.