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

Problem : https://school.programmers.co.kr/learn/courses/30/lessons/152996

프로그래머스

Status : Solved 

Time : 00:54:26

 

---

 

문제 설명

 

어느 공원 놀이터에는 시소가 하나 설치되어 있습니다. 이 시소는 중심으로부터 2(m), 3(m), 4(m) 거리의 지점에 좌석이 하나씩 있습니다.
이 시소를 두 명이 마주 보고 탄다고 할 때, 시소가 평형인 상태에서 각각에 의해 시소에 걸리는 토크의 크기가 서로 상쇄되어 완전한 균형을 이룰 수 있다면 그 두 사람을 시소 짝꿍이라고 합니다. 즉, 탑승한 사람의 무게와 시소 축과 좌석 간의 거리의 곱이 양쪽 다 같다면 시소 짝꿍이라고 할 수 있습니다.
사람들의 몸무게 목록 weights이 주어질 때, 시소 짝꿍이 몇 쌍 존재하는지 구하여 return 하도록 solution 함수를 완성해주세요.

 

입력 및 출력

 

제한사항

• 2 ≤ weights의 길이 ≤ 100,000
• 100 ≤ weights[i] ≤ 1,000
  • 몸무게 단위는 N(뉴턴)으로 주어집니다.
  • 몸무게는 모두 정수입니다.

 

입출력 예

weights	result
[100,180,360,100,270]	4

 

 

---

 

풀이

문제가 어려운 게 아니었는데도, 처음에 개념을 잘못 잡아서 해매버렸다...! 결국 질문하기로 힌트를 조금 얻고 나서야 바로 풀 수 있었던 점이 아쉬울 따름이다.

 

weights 리스트는 최대 100,000개의 길이를 가질 수 있으므로, 단순히 이중 for문으로 접근하면 최대 O(N^2) = 100억의 연산을 가해야 하므로 시간초과를 받게 된다. 하지만 weights 원소의 크기는 100 이상 1000 이하이며,  앉을 수 있는 시소의 길이는 [2, 3, 4]로 정해져 있다. 이말인즉슨, weights의 원소를 기준으로 문제를 풀이해야 한다는 의미이다. 즉 딕셔너리 혹은 리스트를 weight 원소를 기준으로 사용하면 연산 크기가 대폭 줄어든다. 

 

좀 더 구체적으로, 시소에서 weights의 한 원소가 다른 원소와 균형을 이루는 경우는 다음 두 가지가 있다.

• 무게가 같은 두 원소 : 이 경우 원소의 개수가 2개 이상일 때 해당되며, n개의 원소 중에서 무작위로 2개를 뽑는 경우의 수이므로 nC2 = n * (n-1) // 2 로 구할 수 있다.
• 무게가 다른 두 원소 : 앉을 수 있는 시소의 길이가 2, 3, 4m로 정해져 있고, 토크는 (무게) x (길이)로 도출되므로 균형을 이룰 수 있는 다른 무게는 (현재 원소 크기) x (길이비)로 구할 수 있다. 시소 짝궁에서 순서를 따지지 않으므로, 길이비가 1 초과인 경우만 따져보아도 된다. 길이비는 [ 3/2, 4/3, 2 ] 의 세 가지가 나올 수 있고, 이 상황에서 경우의 수는 (현재 무게 원소 개수) x ( 다른 무게 원소 개수) 가 된다.

이것을 반영하여 코드를 작성하면 다음과 같아진다.

 

 

풀이 코드

from collections import defaultdict

def solution(weights):
    answer = 0
    length = len(weights)
    weight_dict = defaultdict(int)
    
    for weight in weights :
        weight_dict[weight] += 1
    
    for key, val in weight_dict.items() :
        if val > 1 :
            answer += val * ( val - 1) // 2
        if key*2 in weight_dict :
            answer += val * weight_dict[key*2]
        if key*3 % 2 == 0 and key*3 // 2 in weight_dict :
            answer += val * weight_dict[key*3 // 2]
        if key*4 % 3 == 0 and key*4 // 3 in weight_dict :
            answer += val * weight_dict[key*4 // 3]

    return answer

풀이 완료!