문제

link

나의 풀이

from itertools import combinations

def solution(weights):
    answer = 0
    weight_map = dict()
    for weight in weights:
        if weight in weight_map.keys():
            weight_map[weight] += 1
        else:
            weight_map[weight] = 1
    for key in weight_map.keys():
        count = weight_map[key]
        if count > 1:
            answer += count * (count - 1) // 2
    cases = list(combinations(list(weight_map.keys()), 2))
    for case in cases:
        x, y = case
        value = (weight_map[x] * weight_map[y])
        if (2*x) == (3*y):
            answer += value
            continue
        elif (2*x) == (4*y):
            answer += value
            continue
        elif (3*x) == (2*y):
            answer += value
            continue
        elif (3*x) == (4*y):
            answer += value
            continue
        elif (4*x) == (2*y):
            answer += value
            continue
        elif (4*x) == (3*y):
            answer += value
            continue
    return answer

고찰

이 문제의 경우, 경우의 수를 줄이는 방법이 중요했다.
이때, 자주 사용할 수 있는 패턴은 바로 Counter와 같은 것들을 사용해서 키-값 딕셔너리를 만드는 것이다. 이와 같이 코드를 짜면, 중복된 원소들을 O(k)와 같이 상수 시간 복잡도로 처리 가능하다.

아래는 GPT가 만들어준 코드이다.

from collections import Counter

def solution(weights):
    # 몸무게 카운트
    weight_count = Counter(weights)
    pairs = 0

    # 거리 비율
    ratios = [2/3, 2/4, 3/4]

    # w1에 대해 탐색
    for w1 in weight_count:
        # 같은 몸무게 처리
        if weight_count[w1] > 1:
            pairs += weight_count[w1] * (weight_count[w1] - 1) // 2

        # 다른 몸무게 처리
        for ratio in ratios:
            w2 = w1 * ratio
            if w2 in weight_count:
                pairs += weight_count[w1] * weight_count[w2]

    return pairs

여기서 배울 점은 아래와 같다.
1. Counter 라이브러리 사용해서 간단하게 중복값 셀 수 있는 딕셔너리 생성
2. ratio같은 것들을 이용해서 수학적 공식을 간단하게 생성