[python][프로그래머스] 시소 짝꿍

장선규·2023년 4월 11일
1

알고리즘

목록 보기
38/40

문제링크
https://school.programmers.co.kr/learn/courses/30/lessons/152996

접근1 (시간초과)

weights를 순회하면서, 각 무게마다 무게*2, 무게*3, 무게*4 한 값끼리 모아준다. 그리고 모인 사람들 끼리 중복 없이 조합을 짜면 된다고 생각했다.

풀이1 (시간초과)

weights를 순회하면서, 각 무게마다 무게*2, 무게*3, 무게*4 한 값끼리 모아준다.

	pair = {}

    # 무게에 대하여 모아주기
    for i in range(len(weights)):
        w2, w3, w4 = weights[i] * 2, weights[i] * 3, weights[i] * 4
        pair.setdefault(w2, tuple())
        pair.setdefault(w3, tuple())
        pair.setdefault(w4, tuple())

        pair[w2] += (i,)
        pair[w3] += (i,)
        pair[w4] += (i,)

무게*2, 무게*3, 무게*4를 key로 가지는 딕셔너리pair에 그 사람의 인덱스를 추가한다. 그러면 해당 무게를 만들 수 있는 사람들끼리 모이게 된다.


그 후, 아까 완성시킨 딕셔너리 pair를 돌면서 길이가 2 이상인 집합을 뽑아낸다. 해당하는 무게를 혼자밖에 만들 수 없다면 짝궁을 만들 수 없기 때문에 제외한다.

 	t_set, ans_set = set(), set()  # 짝궁 가능 집합, 실제 쌍

    # 짝궁 가능한 집합 구하기 (길이 2 이상)
    for t in pair.values():  # 30만
        if len(t) >= 2:
            t_set.add(t)

이제 실제 짝꿍의 쌍을 구한다.

	# 실제 쌍 구하기 (시간초과)
    for t in t_set:
        for i in range(len(t) - 1):
            for j in range(i + 1, len(t)):
                ans_set.add((t[i], t[j]))
    return len(ans_set)

근데 여기서 문제가 생긴다...
t_set의 길이(나올 수 있는 무게의 개수)는 최대 4000개 정도이고, 각 무게에 대해 짝꿍이 가능한 집합(여기선 튜플)의 길이는 최대 300개정도 인 것 같다(평균 150개 정도 나옴).
그렇게 되면, 위의 알고리즘으로는 최악의 경우 약 3000*300*300번 정도 수행하여 2억 7천번 연산이 이루어진다.
결과는 역시나 시간초과로 실패!

접근2

weights의 길이가 최대 10만이기 때문에 O(n^2)보다 좋은 알고리즘을 생각해야한다.weights를 한번만 순회하여 답을 구하는 것이 목표!

생각해보면, 두 사람 a,b가 있을 때 이 둘의 몸무게 비율이 1:1, 2:3, 2:4, 3:4 비율을 가지면 짝꿍이 될 수 있다. 이 아이디어를 기반으로 문제를 풀 수 있었다.

  • 참고) 비율이 반대인 경우(3:2, 4:2, 4:3)은 a,b의 순서만 바꾸면 되니까 고려 x. 어차피 다 체크 해볼 것이다.

풀이2

먼저 무게의 비율이 1:1, 즉 같은 무게인 사람을 구한다.

collections에 있는 Counter를 사용하여 해당 무게를 가진 사람 수를 구한다. weight = [100,180,360,100,270]인 경우엔 100이 2개, 나머지는 1개씩이 될 것이다.

어떤 무게 w를 가진 사람이 한명 뿐이라면, 그 사람 뿐이므로 제외한다. 2명 이상이라면 사람 수(여기선 value값이라 v)에서 2명씩 뽑는 경우의 수 만큼 짝꿍을 만들 수 있다.
쉽게 말해 조합 vC2만큼 만들 수 있다.

from collections import Counter

def solution(weights):
    answer = 0
	counter = Counter(weights)
    
	# 1:1
    for k,v in counter.items():
        if v>=2:
            answer+= v*(v-1)//2

만들 수 있는 짝꿍의 수만큼 answer에 더해준다.


그 다음은 2:3, 2:4, 3:4의 무게 비율을 가진 사람을 찾...기전에 앞서 비율이 1:1인 사람들은 체크를 했으니까 동일한 무게를 가진 사람들은 다 지워준다.
어차피 counter에 그 무게에 해당하는 사람들의 수가 있으므로 괜찮다.

weights = set(weights)  # 중복 제거

이제 weights안에 있는 무게 w에 대해 비율이 2:3, 2:4, 3:4인 무게를 찾는다.
찾았으면, 두 무게를 가지는 사람 수를 곱해주고, 이를 answer에 더해준다.

	# 2:3 2:4 3:4 비율 가지면 짝궁 가능함
    for w in weights:
        if w*2/3 in weights:
            answer+= counter[w*2/3] * counter[w]
        if w*2/4 in weights:
            answer+= counter[w*2/4] * counter[w]
        if w*3/4 in weights:
            answer+= counter[w*3/4] * counter[w]
  • 코드에선 w*2/3 이런식으로 되어있어 사실은 3:2 비율을 찾는 것이지만, 상관없다. 3:2 비율을 가진 녀석이 있다는 것은, 2:3 비율을 가진 녀석이 있다는 것과 같기 때문이다.
  • 순서만 바뀐 것이고, 어차피 모든 w에 대해 순회하므로 괜찮다.

정답코드


from collections import Counter


def solution(weights):
    answer = 0
    
    # 1:1
    counter = Counter(weights)
    for k,v in counter.items():
        if v>=2:
            answer+= v*(v-1)//2
    weights = set(weights) 
    
    # 2:3 2:4 3:4 비율 가지면 짝궁 가능함
    for w in weights:
        if w*2/3 in weights:
            answer+= counter[w*2/3] * counter[w]
        if w*2/4 in weights:
            answer+= counter[w*2/4] * counter[w]
        if w*3/4 in weights:
            answer+= counter[w*3/4] * counter[w]
    return answer

profile
코딩연습

0개의 댓글