C++로 구현하는 사전순 K번째로 작은 수 찾기

문제 개요

두 정수 n과 k가 주어졌을 때, 1부터 n까지의 범위에 있는 숫자들을 사전순(lexicographical order)으로 나열했을 때 k번째로 작은 정수를 찾아야 합니다.

예를 들어 n = 14, k = 3이 입력으로 주어진다면, 숫자들을 사전순으로 정렬한 결과는 [1, 10, 11, 12, 13, 14, 2, 3, 4, 5, 6, 7, 8, 9]와 같습니다. 따라서 세 번째 숫자인 11이 출력됩니다.

접근 방법

사전순 정렬에서 숫자들은 마치 10진 트리(denary tree)처럼 배치됩니다. 즉, 각 숫자는 자신 뒤에 0부터 9까지의 숫자를 붙인 자식 노드를 가지는 구조입니다. 이 구조를 활용하면 현재 위치에서 다음 숫자로 이동할 때 필요한 단계 수를 계산하여 효율적으로 답을 찾을 수 있습니다.

핵심 아이디어는 다음과 같습니다.

findKthNumber() 함수를 정의합니다. 이 함수는 n과 k를 매개변수로 받습니다.
curr := 1로 초기화하고, k를 1 감소시킵니다.
k가 0이 아닌 동안 다음 과정을 반복합니다.
- steps := calcSteps(n, curr, curr + 1)을 호출하여 현재 접두사 구간의 단계 수를 계산합니다.
- steps <= k라면, 현재 서브트리 전체를 건너뛸 수 있으므로 k에서 steps를 빼고 curr을 1 증가시킵니다.
- 그렇지 않다면, curr := curr * 10으로 한 자리 내려가고 k를 1 감소시킵니다.
curr을 반환합니다.

calcSteps() 함수는 nax, n1, n2를 매개변수로 받으며, 특정 접두사 구간에 속하는 숫자의 개수를 계산합니다.

ret := 0으로 초기화합니다.
n1 <= nax인 동안 다음 과정을 반복합니다.
- ret에 min(nax + 1, n2) − n1을 더합니다.
- n1 := n1 * 10, n2 := n2 * 10으로 갱신합니다.
ret을 반환합니다.

구현 예제

아래 코드를 통해 더 자세히 이해해 보겠습니다.

#include <bits/stdc++.h>
using namespace std;
typedef long long int lli;
class Solution {
public:
   int findKthNumber(int n, int k) {
      int curr = 1;
      k--;
      while(k){
         int steps = calcSteps(n, curr, curr + 1);
         if(steps <= k){
            k -= steps;
            curr++;
         }else{
            curr *= 10;
            k -= 1;
         }
      }
      return curr;
   }
   int calcSteps(lli nax, lli n1, lli n2){
      int ret = 0;
      while(n1 <= nax){
         ret += min(nax + 1, n2) - n1;
         n1 *= 10;
         n2 *= 10;
      }
      return ret;
   }
};
main(){
   Solution ob;
   cout << (ob.findKthNumber(14,3));
}

입력

14,3