[프로그래머스] 시소 짝꿍(python)

문제

문제 설명

• 시소의 거리는 2m, 3m, 4m로 정의되어 있다.
• 문제에서 평형이란 두 수 n1, n2와 시소의 거리 d1, d2가 있을 때 n1*d1 = n2*d2를 만족하는 것
• 숫자가 담긴 배열이 주어질 때 여기서 평형을 이루는 쌍의 개수는 몇 개인지 구하는 것(순서 고려x)

풀이

• 배열에 중복되는 수가 존재할 수 있기 때문에 dictionary로 개수를 관리하고, set을 형성
• set의 원소들을 순회하며 계산
• 계산의 편의성을 위해 숫자를 정렬하여, 비례 값으로 나올 수 있는 경우를 리스트로 만들어 확인하였다.

def solution(weights):
    weight_dict = {}
    for weight in weights:
        weight_dict[weight] = weight_dict.get(weight, 0) + 1
    weight_lst = sorted(weight_dict.keys())
    l = len(weight_lst)
    result = 0
    for i in range(l-1):
        for j in range(i+1, l):
            w1 = weight_lst[i]
            w2 = weight_lst[j]
            if w2/w1 in [1/1, 2/1, 4/3, 3/2]:
                result += weight_dict[w1] * weight_dict[w2]
    for k, v in weight_dict.items():
        if v>=2:
            result += (v * (v-1))//2
    return result