큰 수 만들기 (lv2)

문제 설명

어떤 숫자에서 k개의 수를 제거했을 때 얻을 수 있는 가장 큰 숫자를 구하려 합니다.

예를 들어, 숫자 1924에서 수 두 개를 제거하면 [19, 12, 14, 92, 94, 24] 를 만들 수 있습니다. 이 중 가장 큰 숫자는 94 입니다.

문자열 형식으로 숫자 number와 제거할 수의 개수 k가 solution 함수의 매개변수로 주어집니다. number에서 k 개의 수를 제거했을 때 만들 수 있는 수 중 가장 큰 숫자를 문자열 형태로 return 하도록 solution 함수를 완성하세요.

제한 조건

• number는 2자리 이상, 1,000,000자리 이하인 숫자입니다.
• k는 1 이상 number의 자릿수 미만인 자연수입니다.

---

 상당히 재밌는 문제다. 가장 먼저, 처음에 떠올렸던 아이디어를 구현하는 과정을 적어보도록 하겠다. 그리고 나서 좀 더 '그리디'스럽고 깔끔한 풀이를 가져와 비교해보도록 하겠다.

 처음에 떠올린 아이디어는, 문제에서 결국 구해야 하는 문자열의 길이가 len(number)-k 이니까 그 길이 만큼 반복문을 돌리며 문자열 number의 k+1+i 까지의 범위 중 가장 큰 수를 골라 answer에 붙여주는 것이었다. 예를 들면 "1924" 가 number로 주어졌다고 하자. k는 2이므로, answer의 길이는 4 - 2 = 2 가 될 것이다. 따라서 2에서 1만큼 뺀 길이 만큼만 남겨둔 "192" 중 가장 큰 값을 골라 answer에 '9'를 붙여준 뒤, '9' 보다 뒤에있는 "24" 중 가장 큰 값을 골라 다시 answer에 붙여 최종적으로 "94"를 얻는다. 이 과정은 문자열을 slice한 뒤 max를 사용하므로, 결국 O(n^2) 정도의 시간복잡도로 해당 문제를 해결하기 어렵다.

TimeComplexity - Python Wiki

TimeComplexity - Python Wiki

 따라서 일일히 max를 이용해 최대값을 찾는것은 포기하고, 대신 heapq를 이용해 최대값을 가져오는 방법을 선택했다. heapq의 heappop 메서드는 O(log(n)) 복잡도를 가지고 있으므로 slice와 max를 쓰는 방법보다는 실행속도가 여유로워 가까스로 테스트 케이스들을 통과할 수 있었다...

코드

import heapq

def solution(number, k):
    answer = ''
    prev_idx = -1
    heap = [(-int(v),i) for i,v in enumerate(number[0:k+1])]
    heapq.heapify(heap)
    for i in range(len(number)-k):
        while heap:
            num,idx = heapq.heappop(heap)
            if idx > prev_idx:
                answer += str(-num)
                prev_idx = idx
                break
        try:
            heapq.heappush(heap,(-int(number[k+1+i]),k+1+i))
        except IndexError:
            pass
    return answer

 

 하지만 문제를 조금 다르게 접근해 상기했던 대로 좀 더 "그리디"하게 문제를 해결할 수 있다. 주어진 number 배열에서 연속되는 부분 배열이 최대값이 되려면 어떻게 해야겠는가? 처음 생각했던 대로, 주어진 조건처럼 전체 문자열에서 k개의 문자만 제외하며 앞에서부터 가장 큰 값을 골라가는 것이다. 예를 들어 "4177252841" 와 k=4가 주어졌다고 하자. 앞에 두개 정도는 날릴 수 있으므로, 첫 번째 수로 세 번째 자리에 있는 7을 선택하는것이 최적일 것이다.

• "4177252841"
• "77252841"

 이제 수 두개를 날릴 수 있다. 남아있는 문자열에서 세 번째 자리가 2 인것을 확인 가능한데, 2 뒤의 문자들이 2보다 크므로 2를 날려주는게 최선일 것이다.

• "7752841"

 k가 1만큼 남아 있다. 네 번째 자리에 2가 보이는데, 2 뒤에 2 보다 큰 8이 있으므로 2를 제거하는것이 합리적일 것이다. 최종적으로 다음과 같은 해를 얻는다.

• "775841"

 이 예제를 통해 따라간 과정은 마치 스택과 같이 동작하고 있다. 주어진 number를 순회하며 각 숫자를 stack에 넣어주고, 날려야 하는 숫자의 개수가 0이 되지 않는 한 "큰 수 만들기" 에서는 우선적으로 앞의 숫자들이 큰 수가 되어야 할 필요가 있으므로 stack 내에서 현 시점의 수 보다 작은 숫자들은 모두 꺼내버리는 것이다. 다시말하면 answer을 만드는데 최우선적으로 앞자리 수일수록 큰 수가 오도록 한다는 것이다. 이 풀이를 생각한 사람은 이 문제에서 스택 구조를 잘도 생각해낸것 같다...

코드

def solution(number, k):
    stack = [number[0]]
    for num in number[1:]:
        while len(stack) > 0 and stack[-1] < num and k > 0:
            k -= 1
            stack.pop()
        stack.append(num)
    if k != 0:
        stack = stack[:-k]
    return ''.join(stack)