시소 짝꿍 - 파이썬(Python)

문제 링크

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

프로그래머스

---

문제 설명

어느 공원 놀이터에는 시소가 하나 설치되어 있습니다. 이 시소는 중심으로부터 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

---

입출력 예 설명

{100, 100} 은 서로 같은 거리에 마주보고 앉으면 균형을 이룹니다.
{180, 360} 은 각각 4(m), 2(m) 거리에 마주보고 앉으면 균형을 이룹니다.
{180, 270} 은 각각 3(m), 2(m) 거리에 마주보고 앉으면 균형을 이룹니다.
{270, 360} 은 각각 4(m), 3(m) 거리에 마주보고 앉으면 균형을 이룹니다.

---

풀이 과정

• 완전 탐색이라 생각하고 풀었다.
• weights의 길이를 보니 중복되는 부분을 쳐내고 해야 함을 간접적으로 알게 된다.
  • 방법은? Counter사용
• 경우의 수만 구하는 거라면 itertools쓰는 것보다 math.factorial 활용하는 것이 더 빠르다. (itertools 쓰면 시간 초과)
  • 리스트 생성하고 하면 오버헤드 커서 당연함

정답

from itertools import combinations
from collections import defaultdict, Counter
import math
from itertools import combinations
def solution(weights):
    answer = 0
 
    cases = [(1, 1), (2, 3), (1, 2), (3, 2), (2, 1), (3, 4), (4, 3)]
    
    counter = Counter(weights)
    
    # 중복되는 부분 끼리 계산 -> ex. 100 & 100
    for weight, num in counter.most_common():
        if num >= 2:
            answer += math.factorial(num) / (math.factorial(num-2)* math.factorial(2))
        else:
            break
    
    # 중복되는 부분 빼고 계산 후 경우의 수 계산
    for w1, w2 in combinations(counter.keys(), 2):
        for x, y in cases:
            if w1*x == w2*y:
                answer += counter[w1]*counter[w2]
    
                break
    
    return answer