[파이썬] 프로그래머스 : 시소 짝꿍 (Lv.2)

[파이썬] 프로그래머스 : 시소 짝꿍 (Lv.2)

 

프로그래머스

 

---

풀이

0. 방향성 생각

집합 / 딕셔너리 이용해서 시간 줄이기.

일단 케이스가 1:1 비율, 몸무게가 같아서 짝을 이루는 경우가 있다. 이 때는 nC2 = n*(n-1)//2로 짝의 수를 구한다.

비율이 1:1이 아닌 경우에는 따로 구한다.

 

1. 입력 받기

def solution(weights):
    
    table,point = {},{} # 쌍 a,b는 어떤 point에서 균형을 이룬다. point를 만드는 후보 전부 넣기
    for i in weights: # table에는 사람 몸무게 업데이트
        if i not in table:
            table[i] = 1
        else :
            table[i] += 1
        for ni in (2*i,3*i,4*i): # point 짝이 맞는 사람들 넣기
            if ni not in point:
                point[ni] = set()
            point[ni].add(i)

table에는 각 몸무게의 사람이 몇 명 있는지 계산해준다.

point는 몸무게가 i인 사람이 토크를 받아서 몸무게가 ni가 됐을 때 그 몸무게 i를 저장하는 딕셔너리이다.

즉 키 X : 값 set() 에서 set에는 X가 될 수 있는 후보군을 모두 저장한다.

 

밑에는 각각 table, point를 출력한 결과

입력값 〉	[100, 180, 360, 100, 270]
기댓값 〉	4
실행 결과 〉	테스트를 통과하였습니다.
출력 〉	{100: 2, 180: 1, 360: 1, 270: 1}
{200: {100}, 300: {100}, 400: {100}, 360: {180}, 540: {180, 270}, 720: {360, 180}, 1080: {360, 270}, 1440: {360}, 810: {270}}

샘플 테스트 케이스

입력값 〉	[300, 300, 400, 400, 400, 600]
기댓값 〉	15
실행 결과 〉	테스트를 통과하였습니다.
출력 〉	{300: 2, 400: 3, 600: 1}
{600: {300}, 900: {300}, 1200: {400, 600, 300}, 800: {400}, 1600: {400}, 1800: {600}, 2400: {600}}

추가한 테스트 케이스

 

2. 정답 계산하기

    answer = 0
    for x in table.keys():
        if table[x] >= 2:
            answer += table[x]*(table[x]-1)//2

    for x in point.keys():
        if len(point[x]) >= 2:
            temp = []
            for nx in point[x]:
                temp.append(table[nx])
            for i in range(len(temp)):
                for j in range(i+1,len(temp)):
                    answer += temp[i]*temp[j]
    return answer

table에서 몸무게가 같은 경우에는(table[x]>=2 인 경우) 그 집단에서 2명을 뽑는 수가 정답이다.

point에서는 몸무게가 다른 경우에 비율이 맞춰진다.

point[x]의 길이가 2 이상인 경우 서로 다른 몸무게가 비율을 맞춰서 몸무게 x에 도달한다.

 

이 때 그 후보군에 대해서 몇 번 나왔는지 temp에 저장해준다.

각 원소끼리 모두 곱해주면 된다. temp 리스트에서 2개 뽑아서 곱한 후 answer에 더해주면 된다.

코드에서는 2중 for문으로 조합 구현.

---

전체코드

def solution(weights):
    
    table,point = {},{} # 쌍 a,b는 어떤 point에서 균형을 이룬다. point를 만드는 후보 전부 넣기
    for i in weights: # table에는 사람 몸무게 업데이트
        if i not in table:
            table[i] = 1
        else :
            table[i] += 1
        for ni in (2*i,3*i,4*i): # point 짝이 맞는 사람들 넣기
            if ni not in point:
                point[ni] = set()
            point[ni].add(i)
                
    answer = 0
    for x in table.keys():
        if table[x] >= 2:
            answer += table[x]*(table[x]-1)//2

    for x in point.keys():
        if len(point[x]) >= 2:
            temp = []
            for nx in point[x]:
                temp.append(table[nx])
            for i in range(len(temp)):
                for j in range(i+1,len(temp)):
                    answer += temp[i]*temp[j]
    return answer

 

코멘트

오늘 문제 너무잘풀린다.

입력 크기가 10*5라서 일반적인 N^2 방법으로는 안될거같아서 다른 방법 찾다보면 어렵지 않게 풀 수 있다고 생각.