프로그래머스|시소 짝꿍

문제 설명

시소에 세 자리가 있는데 이 자리는 시소의 중앙에서 각각 2, 3, 4m 떨어져 있습니다. 시소에 가해지는 무게는 각 자리의 거리*타는 사람의 무게일 때 두 사람이 시소가 서로 평형을 이루도록 앉을 수 있는 경우의 수를 구하는 문제입니다.

작동 순서

1. 사람의 수가 10만명까지이기 때문에 모든 경우를 다 비교할 경우 시간초과가 발생하므로 다른 방법을 찾아야 하는데 무게는 100에서 1000까지 밖에 없기 때문에 무게별로 사람 수를 저장합니다.
2. 각 무게 별로 무게가 같은 경우(서로 [2,2], [3,3], [4,4]에 앉은 경우), 무게가 3/2배인 경우([2, 3]에 앉은 경우), 무게가 2배인 경우([2,4]에 앉은 경우), 무게가 4/3배인 경우([3,4]에 앉은 경우)의 경우의 수를 구해줍니다. 짝의 수를 구하는 문제이므로 서로 앉은 위치는 상관없이 앉은 사람만 같으면 됩니다.
3. 무게가 같은 경우 n개 중 2개를 고르는 경우와 같으므로 ((무게가 i인 사람의 수)*(무게가 i인 사람의 수-1)/2)를 answer에 더해줍니다.
4. 무게가 각각 3/2배, 2배, 4/3배인 경우 (무게가 i인 사람의 수)(무게가 i(3/2, 2, 4/3 중 하나)인 사람의 수)를 answer에 더해줍니다.
5. 모든 연산이 끝나면 지금까지 구한 경우의 수의 개수를 출력합니다.

소스코드

def solution(weights):
    answer = 0
    arr=[0]*1001
    for i in weights:
        arr[i]+=1
    
    for i in range(1001):
        for j in [1, 3/2, 2, 4/3]:
            num=int(i*j)
            if i*j!=num:
                continue
            if num<1001:
                if num==i:
                    if arr[i]>1:
                        answer+=(arr[i]*(arr[i]-1))//2
                else:
                    answer+=arr[i]*arr[num]
    return answer