[LeetCode] 402 | Remove K Digits | Python

---

문제 링크: https://leetcode.com/problems/remove-k-digits/

---

문제

Given string num representing a non-negative integer num, and an integer k, return the smallest possible integer after removing k digits from num.

---

조건

• 1 <= k <= num.length <= 105
• num consists of only digits.
• num does not have any leading zeros except for the zero itself.

---

내 풀이

맨 앞에서부터 연속된 숫자를 비교해 앞자리가 더 클 때, 앞자리는 버리는 방식이다.

스택(ans)에 num의 맨 첫자리를 넣는다.

그 후, ans의 마지막 자리와 num의 현재자리의 수(digit)를 비교했을 때 digit이 더 크다면 앞자리보다 뒷자리가 더 크므로 이상적인 상황이다. 그렇기 때문에 ans에 추가한다.

만약 digit이 더 작다면 앞자리가 더 큰 상황이기 때문에 ans의 맨 마지막 자리를 pop한다.

또한 이 상황에서는 앞자리를 버리기 때문에 k -= 1을 진행한다.

이를 digit이 더 작지 않을 때까지 반복한다.

 

위 방식을 num의 모든 자리에 대해 반복한 후에도 k가 0이 아니라면 맨 뒤부터 하나씩 차례로 pop한다.

---

코드

class Solution:
    def removeKdigits(self, num: str, k: int) -> str:
        if len(num) == k:   return '0'
        
        ans = [num[0]]
        for digit in num[1:]:
            digit = int(digit)
            
            if k == 0 or digit > int(ans[-1]):
                ans.append(str(digit))
            else:
                while len(ans) > 0 and k > 0 and digit < int(ans[-1]):
                    ans.pop()
                    k -= 1
                ans.append(str(digit))
            
        while k > 0:
            ans.pop()
            k -= 1
        
        ans = ''.join(ans).lstrip('0')
        if ans == '':
            return '0'
        
        return ans