字节题库——440. 字典序的第K小数字 树上递归优化

经典的树上前缀和找第k个节点。

比如权值线段树就用到了这个技巧。

如上图，这颗树的先序遍历就是字典序从小到大。

现在要求先序遍历第k个节点是多少。

我们维护一个函数cal(x)，求的是：x节点的子树包含的节点个数。

显然对于一个节点x，根据上述树的性质观察可知：

其下一层子节点的值范围：[x*10,min(n,2*10 - 1)]

下下一层的节点值范围为：[x*100,min(n,2*100 - 1)]

于是cal函数就可写了。

然后就是经典的dfs找第k大。

若某个节点的子树包含节点数大于k，则往其子树找，否则往其右兄弟节点找(对应二叉树的右子树)。

typedef long long ll;
class Solution {
public:
/*
1
1       2
1 2     1 2

*/
    ll cal(ll pre,ll n){
        ll nxt=pre+1;
        ll ans=0;
        while(pre<=n){
            ans+=min(n+1,nxt)-pre;
            pre*=10;nxt*=10;
        }
        return ans;
    }
    int findKthNumber(int n, int k) {
        int x=1;
        while(1){
            if(k==1){
                return x;
            }
            int nm=cal(x,n);
            if(nm>=k){
                x*=10;
                k--;
            }else{
                x++;
                k-=nm;
            }
        }
        return x;
    }
};