K-th Smallest in Lexicographical Order

KthSmallestInLexicographicalOrder.java

public class Solution {
    	public int findKthNumber(int n, int k) {
		int len = String.valueOf(n).length();
		int [] array = new int [len];
		array[0] = 1;
		for (int i = 1; i < len; i++) {
			array[i] = array[i - 1] * 10 + 1; 
		}
		
		Stack<Integer> s = new Stack<>();
		for (int i = Math.min(9, n); i >= 1; i--) {
			s.push(i);
		}
		while (k > 1) {
			int i = s.pop();
			if (canRemoveSubtree(n, i) && subtreeCount(array, i) < k) {
//				System.out.println("i : " + i + " / k : " + k +"/ subtreeCount(array, i):" + subtreeCount(array, i));
				k -= subtreeCount(array, i);
				continue;
			}
			
			k--;
			for (long j = i * 10 + 9; j >= i * 10L; j--) {
				if (j <= n) {
					s.push((int)j);
				}
			}
		}
		return s.pop();
	}
	
	private int subtreeCount(int[] array, int i) {
		return array[array.length - String.valueOf(i).length()];
	}
	
	private boolean canRemoveSubtree(int n, int i) {
		int len = String.valueOf(n).length();
		if (len <= String.valueOf(i).length()) {
			return false;
		}
		while (len != String.valueOf(i).length()) {
			i *= 10;
			i += 9;
		}
		return i <= n;
	}
}