[프로그래머스] 시소 짝꿍 | 파이썬

---

문제설명

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

 

제한사항

2 ≤ weights의 길이 ≤ 100,000

100 ≤ weights[i] ≤ 1,000

• 몸무게 단위는 N(뉴턴)으로 주어집니다.
• 몸무게는 모두 정수입니다.

 

입출력 예

---

문제풀이

• 이번 문제는 시간복잡도에 대해서 다시 한번 고려해 볼 수 있는 시간이었다. 우선 생각났던 풀이는 이중 for문을 통하여 weights의 원소를 중복없이 2개씩 짝을 지어주고 그 두 값의 비율이 시소의 거리와 같은지를 확인하는 것이었다. 하지만 제한 사항에서 weights의 길이의 최댓값이 100,000이기 때문에 이중 for문을 사용하게되면 시간복잡도 O(n^2)으로 10억회의 연산을 수행하기 때문에 무조건 시간초과가 발생한다. 따라서 weights리스트의 접근이 아니라 weights가 가지고있는 원소(weight)에 대하여 접근해야한다.(dictionary 사용이 필요하다.)
• 따라서 defaultdict(참고 : https://dongdongfather.tistory.com/69)을 import하여 weights의 원소들을 dict에 key로 담아주고 해당 원소들의 갯수를 val로 담아준다. 그리고 각 key값을 확인하면서 key를 기준으로 생길 수 있는 시소 짝꿍을 조건으로 확인하면서 세어준다.
• 우선 같은 무게를 가진 원소가 2개 이상이라면 nC2[(n * (n - 1)) // 2]의 조합으로 짝꿍이 생길 수 있으며 그리고 각 key에 대해서 시소의 좌석의 비율이 1 보다 큰 비율들을 확인하며 짝꿍의 수를 추가해준다. 이때 1보다 큰 비율만 확인하면 되는 이유는 모든 원소들을 전부 돌면서 확인을 해주기 때문에 결국 짝꿍을 중복해서 넣어주지 않기 위함이다. 이렇게 딕셔너리의 모든 원소들을 순회하면 요구하는 답을 얻은 수 있다.

코드 & 설명

from collections import defaultdict

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