[SWEA]Ugly Number - 파이썬

---

문제

Ugly Number란 소인수분해 했을 때, 2 / 3 / 5 로만 이루어진 수를 의미한다. 

예외로 숫자 1은 첫번째 Ugly Number로 정의한다.

Ugly Number를 순서대로 나열하면 다음과 같다.

1, 2, 3, 4, 5, 6, 8, 9, 10, 12, ....

이때, 숫자 K를 입력받으면, 'K번째' Ugly Number를 출력하시오.

 

첫줄에는 문제의 개수 Q(1 <= Q <= 10,000)가 입력된다. 

두번째 줄에는 K가 Q개만큼 공백을 두고 입력된다. (1<= K <= 1,500)

각각의 K에 대해 공백을 두고 K번째의 Ugly Number를 출력하시오.

 

[입력]

3
1 9 10

 

[출력]

1 10 12

---

풀이

처음에는 소인수가 2, 3, 5밖에 없는지 확인하는 함수를 만들어서

2부터 함수를 돌려서 해당하면 heap에 집어넣고, 그 힙의 길이가 가장 큰 K와 같아질 때 까지 하려고 했으나

당연히 시간 초과가 나서 한참 고민했었다.

여기저기 찾아본 결과 풀이를 세개정도 찾을 수 있었다.

---

1. DP를 이용한다.

Q = int(input())
arr = list(map(int, input().split()))

ugly_numbers = [1]  # 1은 첫 번째 Ugly Number
next_2, next_3, next_5 = 2, 3, 5  # 다음에 곱할 수
index_2, index_3, index_5 = 0, 0, 0  # 각각의 다음 인덱스

for _ in range(max(arr)):
    next_ugly = min(next_2, next_3, next_5)  # 다음 계산은 작은 것부터
    ugly_numbers.append(next_ugly)

    if next_ugly == next_2:
        index_2 += 1
        next_2 = ugly_numbers[index_2] * 2
    if next_ugly == next_3:
        index_3 += 1
        next_3 = ugly_numbers[index_3] * 3
    if next_ugly == next_5:
        index_5 += 1
        next_5 = ugly_numbers[index_5] * 5
        
ans = []
for q in arr:
    ans.append(ugly_numbers[q - 1])
print(*ans)

다음 숫자를 고를 때, 최대한 작은 숫자를 골라서 넣어서 ugly_numbers 안에 들어가는 순서가 작은 숫자의 순서대로

들어갈 수 있도록 해준다.

 

2. 힙을 이용

import heapq

N = int(input())
k = list(map(int,input().split()))
ugly = []
heap = [1]	
heapq.heappush(heap, 2)
heapq.heappush(heap, 3)
heapq.heappush(heap, 5)

while len(ugly) < max(k):
    n = heapq.heappop(heap)	# 힙에서 숫자를 빼온 후
    if n not in ugly:		# 만약 최종 리스트에 없다면
        ugly.append(n)		# 추가를 해주고
        heapq.heappush(heap, 2 * n)	# 2, 3, 5를 곱해서 각각 넣어준다.
        heapq.heappush(heap, 3 * n)
        heapq.heappush(heap, 5 * n)

for i in k:
    print(ugly[i - 1], end =' ')

힙(heap)이란 완전이진트리의 한 종류로, 키값이 가장 크거나 작은 노드를 찾기 위해 만들어진 자료구조

• 최대 힙(max heap): 부모노드의 키값 > 자식 노드의 키값, 루트노드가 가장 큼
• 최소 힙(min heap): 부모노드의 키값 < 자식 노드의 키값, 루트노드가 가장 작음

위의 풀이에서는 곱한 숫자들을 바로바로 heap에 넣어서 heappop에서 나오는 숫자가 자동으로 가장 작은 숫자가 되도록 해준다.

 

초기 상태에서 ugly num을 담을 비어있는 ugly list와 heap을 생성한다.

heappush를 이용해 heap에 2, 3, 5를 넣어준다. 이때, 힙의 특성으로 인해 숫자 순서대로 들어가게 된다.

while문에서 heappop을 이용해서 가장 작은 원소인 1을 빼준다.

이후 만약 그 1이 ugly에 존재하지 않는다면 ugly에 추가를 해준 후, n에 2, 3, 5를 곱해서 다시 추가한다.

n = 1인 경우에는 heap에 새롭게 추가 되는 건 없지만, 두번째 while문을 돌 때부터는 n = 2가 되면서 2에 2, 3, 5를 곱한 4, 6, 10이 추가 되어서 heap = [3, 4, 5, 6, 10]이 될 것이다.

이후, ugly의 길이가 구하고자하는 길이가 된다면 while문을 멈춰서 답을 구하면 된다. 

 

3. 다중 포인터 이용

Q = int(input())
arr = list(map(int,input().split()))

ugly = [1]
p2, p3, p5 = 0, 0, 0

while len(ugly) <= max(arr):
    next_ugly = min(ugly[p2] * 2, ugly[p3] * 3, ugly[p5] * 5)
    ugly.append(next_ugly)

    if next_ugly == ugly[p2] * 2:
        p2 += 1
    if next_ugly == ugly[p3] * 3:
        p3 += 1
    if next_ugly == ugly[p5] * 5:
        p5 += 1

ans = []
for q in arr:
    ans.append(ugly[q - 1])
print(*ans)

GPT가 알려준 방법으로 DP와 크게 다를 건 없어보인다

p2, p3, p5는 각각 ugly number를 소인수 분해했을 때 2, 3, 5의 지수와 같다. 

이 방법이 제일 간결하고 이해하기 직관적이라 기억해 둘 필요가 있어보인다.

 

---

• min함수를 좀 더 적극 활용하자
• heap의 활용 관련 문제를 좀 더 풀어보자