用于打印/获取二叉搜索树中第 k 个最小元素的 SINGLE 递归函数

Question

我正在尝试打印 BST 中的第 k 个最小元素。 第一个解决方案是使用中序遍历。 下一个解决方案是通过计算其左子树的大小来找到当前节点的索引。 完整的算法：

Find size of left subtree:
   1.If size = k-1, return current node
   2.If size>k return (size-k)th node in right subtree
   3.If size<k return kth node in left subtree

这可以使用一个单独的计数函数来实现，看起来像

public class Solution {
    public int kthSmallest(TreeNode root, int k) {
        //what happens if root == null
        //what happens if k > total size of tree
        return kthSmallestNode(root,k).val;

    }

    public static TreeNode kthSmallestNode(TreeNode root,int k){
        if(root==null) return root;
        int numberOfNodes = countNodes(root.left);

        if(k == numberOfNodes ) return root;
        if(k<numberOfNodes ) return kthSmallestNode(root.left,k);
        else return kthSmallestNode(root.right,k-numberOfNodes );
    } 

    private static int countNodes(TreeNode node){
        if(node == null) return 0;
        else return 1+countNodes(node.left)+countNodes(node.right);
    }
}

但我看到我们多次计算同一棵树的大小，因此一种方法是维护一个数组来存储这些大小，就像 DP 方式一样。

但我想为此编写一个递归解决方案。这是我编写的代码。

class Node {
    int data;
    Node left;
    Node right;

    public Node(int data, Node left, Node right) {
        this.left = left;
        this.data = data;
        this.right = right;
    }
}

public class KthInBST
{
    public static Node createBST(int headData)
    {
          Node head = new Node(headData, null, null);
          //System.out.println(head.data);
          return head;
    }

    public static void insertIntoBst(Node head, int data) 
    {

            Node newNode = new Node(data, null, null);
            while(true) {

            if (data > head.data) {
                if (head.right == null) {
                    head.right = newNode;
                    break;
                } else {
                    head = head.right;
                }
            } else {
                if (head.left == null) {
                    head.left = newNode;
                    break;
                } else {
                    head = head.left;
                }
            }
          }
    }

    public static void main(String[] args) 
    {
           Node head = createBST(5);
           insertIntoBst(head, 7);
           insertIntoBst(head, 6);

           insertIntoBst(head, 2);
           insertIntoBst(head, 1);
           insertIntoBst(head, 21);
           insertIntoBst(head, 11);
           insertIntoBst(head, 14);
           insertIntoBst(head, 3);


           printKthElement(head, 3);
    }

    public static int printKthElement(Node head, int k)
    {
         if (head == null) {
            return 0;
         }

         int leftIndex  = printKthElement(head.left, k);

         int index = leftIndex + 1;


         if (index == k) {
            System.out.println(head.data);
         } else if (k > index) {
            k = k - index;
            printKthElement(head.right, k);
         } else {
            printKthElement(head.left, k);
         }
         return index;
    }


}

这是打印正确答案但多次，我想通了为什么它打印多次但不明白如何避免它。 而且如果我想返回节点而不是仅仅打印我该怎么做？ 谁能帮我解决这个问题？

Answer 1

目的：

递归查找二叉搜索树中第 k 个最小的元素，并返回该元素对应的节点。

观察：

小于当前元素的元素个数就是左子树的大小，所以我们不递归计算它的大小，而是在class Node中引入一个新成员，即lsize，代表左子树的大小当前节点的子树。

解决方案：

在每个节点我们将左子树的大小与k的当前值进行比较：

1. 如果head.lsize + 1 == k：我们答案中的当前节点。
2. 如果head.lsize + 1 > k：左子树中的元素多于k个，即k个最小的元素在左子树中。所以，我们向左走。
3. 如果head.lsize + 1 < k：当前元素以及左子树中的所有元素都小于我们需要查找的第k个元素。因此，我们转到右子树，但也将 k 减少了左子树中元素的数量 + 1（当前元素）。通过从k 中减去它，我们确保我们已经考虑了小于k 并且以当前节点（包括当前节点本身）的左子树为根的元素的数量。

代码：

class Node {
    int data;
    Node left;
    Node right;
    int lsize;

    public Node(int data, Node left, Node right) {
        this.left = left;
        this.data = data;
        this.right = right;
        lsize = 0;
    }
}

---

public static void insertIntoBst(Node head, int data) {

        Node newNode = new Node(data, null, null);
        while (true) {

            if (data > head.data) {
                if (head.right == null) {
                    head.right = newNode;
                    break;
                } else {
                    head = head.right;
                }
            } else {
                head.lsize++; //as we go left, size of left subtree rooted 
                             //at current node will increase, hence the increment.
                if (head.left == null) {
                    head.left = newNode;
                    break;
                } else {
                    head = head.left;
                }
            }
        }
    }

---

    public static Node printKthElement(Node head, int k) {
        if (head == null) {
            return null;
        }

        if (head.lsize + 1 == k) return head;
        else if (head.lsize + 1 > k) return printKthElement(head.left, k);
        return printKthElement(head.right, k - head.lsize - 1);
    }

---

变化：

1. class Node 中引入了新成员lsize。
2. insertIntoBst 中的轻微修改。
3. printKthElement 的重大变化。

角盒：

添加检查以确保k 介于1 和树的大小之间，否则将返回null 节点，从而导致NullPointerException。

到目前为止，这是我尝试过的测试用例。任何建议或更正都非常受欢迎。 :)