K-th Smallest in Lexicographical Order LeetCode solution in C++

1
class Node {
2
 private:
3
  Node *links[10];
4
  bool flag = false;
5

6
 public:
7
  bool containsKey(int digit) {
8
    return links[digit] != nullptr;
9
  }
10
  void put(int digit, Node *node) {
11
    links[digit] = node;
12
  }
13
  Node *get(int digit) {
14
    return links[digit];
15
  }
16
  bool isEnd() {
17
    return flag;
18
  }
19
  void setEnd() {
20
    flag = true;
21
  }
22
};
23

24
class Trie {
25
 private:
26
  Node *root;
27

28
 public:
29
  Trie() {
30
    root = new Node();
31
  }
32
  void insert(int num) {
33
    string str = to_string(num);
34
    Node *node = root;
35
    for (int i = 0; i < str.length(); i++) {
36
      int digit = str[i] - '0';
37
      if (!node->containsKey(digit)) {
38
        node->put(digit, new Node());
39
      }
40
      node = node->get(digit);
41
    }
42
    node->setEnd();
43
    // stack<int> st;
44
    // while(num) {
45
    //     st.push(num%10);
46
    //     num /= 10;
47
    // }
48
    // Node* node = root;
49
    // while(!st.empty()) {
50
    //     int digit = st.top();
51
    //     st.pop();
52
    //     if(!node->containsKey(digit)) {
53
    //         node->put(digit, new Node());
54
    //     }
55
    //     node = node->get(digit);
56
    // }
57
    // node->setEnd();
58
  }
59
  int kthELE(int k) {
60
    int ans = 0;
61
    dfs(root, ans, k, 0);
62
    return ans;
63
  }
64
  void dfs(Node *root, int &ans, int &k, int num) {
65
    if (k == 0) return;
66
    if (root->isEnd()) {
67
      ans = num;
68
      k--;
69
    }
70
    for (int i = 0; i <= 9; i++) {
71
      if (root->containsKey(i)) {
72
        dfs(root->get(i), ans, k, num * 10 + i);
73
      }
74
    }
75
  }
76
};
77

78
class Solution {
79
 public:
80
  int findKthNumber(int n, int k) {
81
    Trie trie;
82
    for (int i = 1; i <= n; i++) trie.insert(i);
83
    return trie.kthELE(k);
84
  }
85
};