시소 짝

프로그래머스(시소 짝)

문제

• weights에 사람의 몸무게가 주어졌을 때, 2 사람을 뽑아 시소에 태우고자한다.
• 시소는 2m, 3m, 4m 거리에 좌석이 있다고 할 때, 시소가 평행하도록 하는 짝의 수를 구하라.

방법

1. 계산의 편의를 위해 오름차순으로 정렬해준다.
2. Counter를 이용해 몸무게와 해당 몸무게를 가진 사람의 수를 dictionary로 생성한다.
3. 1:1 비율 즉, 동일한 무게를 가진 사람쌍의 수를 먼저 구한다.
   $answer = answer + \frac{n(n-1)}{2}$
4. 2m, 3m, 4m의 조합으로 나올 수 있는 비율 $\frac{3}{4}, \frac{2}{3}, \frac{2}{4}$을 가지는 경우에 대해서 나올 수 있는 조합을 구한다.
   $answer = answer + w_{1, cnt} \times w_{2, cnt}$

코드

from collections import Counter

def solution(weights):
    answer = 0
    weights = sorted(weights)

    dict_w = Counter(weights)
    for i in dict_w.values():
        answer += i*(i-1)//2
    
    w = list(dict_w.items())
    n = len(w)

    for i in range(n):
        for j in range(i+1, n):
            k1, v1 = w[i][0], w[i][1]
            k2, v2 = w[j][0], w[j][1]
            
            if k1 / k2 in [3/4, 2/3, 2/4]:
                answer += v1*v2
            
    return answer