劍指Offer——二叉搜索樹的第k個結點

題目描述

給定一棵二叉搜索樹，請找出其中的第k小的結點。例如， （5，3，7，2，4，6，8） 中，按結點數值大小順序第三小結點的值爲4。

題解

#include <iostream>
#include <string>

using namespace std;

struct TreeNode {
    int val;
    struct TreeNode *left;
    struct TreeNode *right;

    TreeNode(int x) :
            val(x), left(NULL), right(NULL) {
    }

    TreeNode() {}
};

TreeNode *newTree() {
    TreeNode *node = new TreeNode;
    int x;
    cin >> x;
    if (!x)node = NULL;
    else {
        node->val = x;
        node->left = newTree();
        node->right = newTree();
    }
    return node;
}

int count = 0;

TreeNode *KthNode(TreeNode *pRoot, int k) {
    if (pRoot) {
        TreeNode *node = KthNode(pRoot->left, k);
        if (node) return node;
        if (++count == k)return pRoot;
        node = KthNode(pRoot->right, k);
        if (node)return node;
    }
    return NULL;
}

int main() {
    ios::sync_with_stdio(false);
    int k;
    cin >> k;
    TreeNode *root = NULL;
    root = newTree();
    cout << KthNode(root, k)->val;
    return 0;
}