문제 링크: 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의 길이가 100,000 이므로 O(N^2)을 넘지 않도록 하는 것이 관건이라는 생각이 들었다. 완전 탐색을 한다면 각각의 weight에 대해 다른 모든 weight와 짝꿍 비교를 해야하는데 이렇게 되면 N개의 weight를 N-1개의 다른 weight와 비교하게 되므로 O(N^2)를 넘어가게 된다. 따라서 다른 방법을 찾아야 한다.
시간복잡도를 줄이는 가장 일반적인 접근법은 정렬과 해시를 활용하는 방법이다. 필자는 처음에 정렬을 이용해 풀 수 있지 않을까 하는 마음에 시도해보았지만 생각보다 잘 되지않아 중도 포기했다.(방법이 있는지는 아직도 잘 모르겠다...)
그래서 해시를 이용한 접근을 해보았는데 생각보다 어렵지 않게 풀 수 있었다. 일단 모든 weights의 값을 key로 해서 dict에 저장하고 value는 같은 key의 개수를 카운트하기로 했다. 이렇게 되면 중복되는 weight에 대한 처리가 간단해진다. 같은 weight의 경우 Combination을 구해주면 되고 2배, 3/2배, 3/4배의 경우에는 dict[weight] dict[weight 2]와 같은 방식으로 가능한 pair를 계산하면 된다.
풀이 코드
from collections import defaultdict
import math
def nCr(n, r):
return math.factorial(n) // (math.factorial(n-r) * math.factorial(r))
def solution(weights):
result = 0
weights_dict = defaultdict(lambda: 0)
for weight in weights:
weights_dict[weight] += 1
for weight in weights_dict.keys():
if weights_dict[weight] > 1:
result += nCr(weights_dict[weight], 2)
if weight * 2 in weights_dict.keys():
result += weights_dict[weight] * weights_dict[weight * 2]
if weight * 3 / 2 in weights_dict.keys():
result += weights_dict[weight] * weights_dict[weight * 3 / 2]
if weight * 4 / 3 in weights_dict.keys():
result += weights_dict[weight] * weights_dict[weight * 4 / 3]
return result