[프로그래머스] 시소 짝궁

코딩테스트 연습 - 시소 짝궁

문제 설명

---

def solution(weights: List[int]) -> int:
    ...

2(m), 3(m), 4(m) 거리의 지점에 좌석이 있는 시소가 평형이 되는 쌍의 갯수를 구하는 문제이다. solution 함수의 인자에 사람들의 몸무게 목록이 주어진다.

제한사항

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

풀이

시소가 평형을 이루는지 확인하려면 두 사람의 무게와 각 좌석의 거리의 곱이 같은지 확인하면 된다. 이를 응용해서 몸무게1 / 몸무게2 == 거리2 / 거리1로 비율을 확인하면 된다. 문제에서 주어진 좌석의 거리는 2, 3, 4이다.

입력으로 주어지는 weights 배열의 최대 길이는 10만이기 때문에 이 배열에 대해 for문을 2중으로 사용하면 시간 초과가 발생한다. 반면 몸무게의 범위는 100 ~ 1000으로 길이가 더 짧다. 따라서 Counter를 사용해 weights의 각 원소에 대한 등장 횟수를 구하고 이를 순회해서 해결했다.

제출 코드

from collections import Counter
from typing import List

def solution(weights: List[int]) -> int:
    answer = 0
    weights.sort()

    weights_counter = Counter(weights)

    ratio = [1, 1 / 2, 2 / 3, 3 / 4]
    keys = list(weights_counter.keys())
    for i in range(len(keys)):
        weight_src = keys[i]
        answer += int((weights_counter[weight_src] * (weights_counter[weight_src] - 1)) / 2)
        for j in range(i + 1, len(keys)):
            weight_dest = keys[j]
            if weight_src / weight_dest in ratio:
                answer += weights_counter[weight_src] * weights_counter[weight_dest]
    return answer

후기

간단한 문제였지만 생각보다 코드를 복잡하게 구현했다. 다른 사람들의 풀이를 보니 훨씬 간단한 방법이 있었다. weight_counter[src]와 weight_counter[dst]를 비교하는 식이었지만, weight_counter[ratio]로 사람 수를 구하면 코드를 훨씬 직관적으로 작성할 수 있었다.