문제
https://school.programmers.co.kr/learn/courses/30/lessons/152996
풀이과정
문제는 굉장히 간단하다.
시소는 중심으로부터 2, 3, 4m 떨어진 곳에 좌석이 있다(양쪽으로)
중심으로부터의 거리 * 사람의 무게 가 양쪽이 같으면 둘이 시소 짝궁이라고 할 수 있다.
주어진 몸무게를 가진 사람들 중 시소 짝꿍이 몇 쌍 존재하는지 구해야한다.
Weight_A Distance_A = Weight_B Distance_B 가 되는 경우의 수를 구하면 된다고 생각했다.
처음에 아주 단순하게, 4중 반복문을 생각했다.
answer = set()
for idx_a, w_a in enumerate(weights):
for d_a in [2, 3, 4]:
for idx_b, w_b in enumerate(weights[idx_a + 1:]):
for d_b in [2, 3, 4]:
if w_a * d_a == w_b * d_b:
answer.add([idx_a, idx_b + idx_a + 1])
answer = len(answer)위의 방식으로 풀이하니 예시 테스트 케이스는 맞았지만 시간이 아주 오래 걸렸다.
weights의 길이가 최대 100000 였으니, 최악의 경우 10000000000 번의 연산이 필요하므로 시간 초과 + 실패가 떴다.
다음으로는 이것 저것 시도하다가 딕셔너리를 사용해보자는 아이디어가 떠올랐다.
풀이과정은 다음과 같다.
1. 원본 weights에 중복된 수가 있는지 찾아서 same_dict1에 weight와 갯수를 저장하고 중복이 없는 unique_weight 배열을 새로 만든다
2. 시소 중심으로부터의 거리(2, 3, 4)를 unique_weight 각 요소에 곱한 값을 이어붙인 multiply_weight 배열을 만든다
3. multiply_weight의 요소를 앞에서부터 하나씩 꺼내서 뒤에 같은 값이 있는지 확인 후, 같은 값이 있다면 same_dict2에 그 값을 key로 추가하고 value는 []로 둔다.
4. same_dict2의 key를 순회하면서 key를 2, 3, 4로 나눈 값이 unique_weight에 있는지 확인 후, 있다면 value에 추가한다 (weight distance를 해서 같아지는 값이 존재한다는 의미이므로)
5. same_dict2의 value를 순회하면서,
5-1. value의 길이가 3보다 크거나 같으면 2개를 뽑는 경우가 1보다 크다는 의미이므로 모든 경우를 구하고, 각 경우 사용되는 weight에 대해서 원본 weights에 존재하는 횟수를 구해서 곱해준다
(ex: value에서 100, 200이 한 쌍으로 존재하는데 원본 배열에서 100이 3번, 200이 2번 있다면 100, 200 쌍은 3 2번 나올 수 있음)
5-2. value의 길이가 2라면, 각 요소가 원본 weights에 존재하는 횟수를 구해서 곱해준다
=> 곱한 값들은 cnt(=answer)에 더해줌
6. 마지막으로 원본 weights에 동일한 값이 있었던 weight에 대해서 중복 갯수 중 2개를 고르는 경우의 수(조합의 수)를 cnt에 더해준다
이렇게 저렇게 해서... 풀긴 풀었는데 코드가 아주 길다는 점이 아쉽다.
의미 있었던 예시 테스트 케이스는 weights = [100, 100, 100, 150, 150, 200, 300] 였다.
이 경우, 결과가 18이 나와야 하는데 오류가 나는 코드에서는 19가 나왔다. (테스트케이스 3-11 실패)
그 부분을 고쳤더니 통과 되었다!
코드
import math
from itertools import combinations
def solution(weights):
# 같은 수 있는지 확인하고 하나만 남기고 제거 O(N)
unique_weight = [weights[0]]
same_dict1 = {} # weights에 같은 수가 있으면, 어떤 수가 몇 번 나오는지 기록하는 딕셔너리
for w in weights[1:]:
if w in unique_weight:
same_dict1[w] = same_dict1.get(w, 1) + 1
else:
unique_weight.append(w)
# print(same_dict1)
# 2, 3, 4를 곱한 값을 이어 붙이기
multiply_weight = [i * 2 for i in unique_weight] + [i * 3 for i in unique_weight] + [i * 4 for i in unique_weight]
# print(multiply_weight)
# 앞에서부터 하나씩 꺼내서 뒤에 같은 값 있는지 확인
i = 0
same_dict2 = {}
while i < len(multiply_weight) - 1:
w = multiply_weight[i]
if w in multiply_weight[i+1:]: # 같은 값이 뒤에 존재하면, same_dict2에 key를 추가하고 value는 빈 리스트로 설정
same_dict2[w] = []
i += 1
for key in same_dict2.keys(): # 같은 값이 존재하는 수에 대해서, 2, 3, 4로 나눈 뒤 unique_weight에 있는 값인지 확인하고, 있으면 same_dict2[key]에 추가한다
if key / 2 in unique_weight:
same_dict2[key].append(int(key/2))
if key / 3 in unique_weight:
same_dict2[key].append(int(key/3))
if key / 4 in unique_weight:
same_dict2[key].append(int(key/4))
cnt = 0
for value in same_dict2.values():
if len(value) >= 3:
# value가 3개 이상이면 2개를 뽑는 경우의 수가 1보다 커진다
# -> combinations를 이용해서 2개 뽑는 경우를 모두 구하고,
# 원본 배열에서 중복이 있었는지 확인하고, 있다면 중복 갯수를 곱해주는 방식으로 계산
combs = list(combinations(value, 2))
tmp = 0
for comb in combs:
tmp += same_dict1.get(comb[0], 1) * same_dict1.get(comb[1], 1) # 중복이 없었다면 1이 곱해지는 것
else:
tmp = same_dict1.get(value[0], 1) * same_dict1.get(value[1], 1) # 여기도 마찬가지로 원본 배열에 중복이 있다면 중복 숫자의 갯수, 없다면 1이 곱해짐
# print(value, tmp)
cnt += tmp
for value in same_dict1.values(): # 처음 weights에서 동일한 값이 있던 수에 대해서 해당 수 중 2개를 고르는 경우의 수를 cnt에 더해준다
cnt += math.comb(value, 2)
# print(same_dict2)
return cnt
weights = [100, 100, 180, 360, 270]
weights = [100, 100, 100, 150, 150, 200, 300]
# weights = [100, 100, 180, 180, 360, 100, 270]
print(solution(weights))다른 사람의 풀이
from itertools import combinations
from collections import Counter
def solution(weights):
cnt = 0
weights = Counter(weights)
for a, b in combinations(weights.keys(), 2): # 서로 다른 무게
if a*2 == b*3 or a*2 == b*4 or a*3 == b*4 or b*2 == a*3 or b*2 == a*4 or b*3 == a*4:
cnt += weights[a] * weights[b]
for v in weights.values(): # 같은 무게
if v > 1:
cnt += sum([i for i in range(1, v)])
return cntcounter와 combinations를 사용한 깔끔한 풀이.
배워야지...
