給定一個二叉搜索樹(BST),找到樹中第 K 小的節點

題目:給定一個二叉搜索樹(BST),找到樹中第 K 小的節點。
出題人:阿里巴巴出題專家:文景/阿里雲 CDN 資深技術專家
參考答案

  • 考察點
  1. 基礎數據結構的理解和編碼能力
  2. 遞歸使用
  • 示例
       5
      / \
     3   6
    / \
   2   4
  /
 1
 

說明:保證輸入的 K 滿足 1<=K<=(節點數目)

解法1:樹相關的題目,第一眼就想到遞歸求解,左右子樹分別遍歷。聯想到二叉搜索樹的性質,root 大於左子樹,小於右子樹,如果左子樹的節點數目等於 K-1,那麼 root 就是結果,否則如果左子樹節點數目小於 K-1,那麼結果必然在右子樹,否則就在左子樹。因此在搜索的時候同時返回節點數目,跟 K 做對比,就能得出結果了。

/**
 * Definition for a binary tree node.
 **/

public class TreeNode {
    int val;
    TreeNode left;
    TreeNode right;
    TreeNode(int x) { val = x; }
}

class Solution {
    private class ResultType {
    
        boolean found;  // 是否找到
        
        int val;  // 節點數目
        ResultType(boolean found, int val) {
            this.found = found;
            this.val = val;
        }
    }

    public int kthSmallest(TreeNode root, int k) {
        return kthSmallestHelper(root, k).val;
    }

    private ResultType kthSmallestHelper(TreeNode root, int k) {
        if (root == null) {
            return new ResultType(false, 0);
        }

        ResultType left = kthSmallestHelper(root.left, k);

        // 左子樹找到,直接返回
        if (left.found) {
            return new ResultType(true, left.val);
        }

        // 左子樹的節點數目 = K-1,結果爲 root 的值
        if (k - left.val == 1) {
            return new ResultType(true, root.val);
        }

        // 右子樹尋找
        ResultType right = kthSmallestHelper(root.right, k - left.val - 1);
        if (right.found) {
            return new ResultType(true, right.val);
        }

        // 沒找到,返回節點總數
        return new ResultType(false, left.val + 1 + right.val);
    }
}

解法2:基於二叉搜索樹的特性,在中序遍歷的結果中,第k個元素就是本題的解。
最差的情況是k節點是bst的最右葉子節點,不過每個節點的遍歷次數最多是1次
遍歷並不是需要全部做完,使用計數的方式,找到第k個元素就可以退出。
下面是go的一個簡單實現。

// BST is binary search tree
type BST struct {
	key, value  int
	left, right *BST
}

func (bst *BST) setLeft(b *BST) {
	bst.left = b
}

func (bst *BST) setRight(b *BST) {
	bst.right = b
}

// count 查找bst第k個節點的值,未找到就返回0
func count(bst *BST, k int) int {
	if k < 1 {
		return 0
	}

	c := 0
	ok, value := countRecursive(bst, &c, k)

	if ok {
		return value
	}

	return 0
}

// countRecurisive 對bst使用中序遍歷
// 用計數方式控制退出遍歷,參數c就是已遍歷節點數
func countRecursive(bst *BST, c *int, k int) (bool, int) {
	if bst.left != nil {
		ok, value := countRecursive(bst.left, c, k)
		if ok {
			return ok, value
		}
	}

	if *c == k-1 {
		return true, bst.value
	}

	*c++

	if bst.right != nil {
		ok, value := countRecursive(bst.right, c, k)
		if ok {
			return ok, value
		}
	}

	return false, 0
}

// 下面是測試代碼,覆蓋了退化的情況和普通bst

func createBST1() *BST {
	b1 := &BST{key: 1, value: 10}
	b2 := &BST{key: 2, value: 20}
	b3 := &BST{key: 3, value: 30}
	b4 := &BST{key: 4, value: 40}
	b5 := &BST{key: 5, value: 50}
	b6 := &BST{key: 6, value: 60}
	b7 := &BST{key: 7, value: 70}
	b8 := &BST{key: 8, value: 80}
	b9 := &BST{key: 9, value: 90}

	b9.setLeft(b8)
	b8.setLeft(b7)
	b7.setLeft(b6)
	b6.setLeft(b5)
	b5.setLeft(b4)
	b4.setLeft(b3)
	b3.setLeft(b2)
	b2.setLeft(b1)

	return b9
}

func createBST2() *BST {
	b1 := &BST{key: 1, value: 10}
	b2 := &BST{key: 2, value: 20}
	b3 := &BST{key: 3, value: 30}
	b4 := &BST{key: 4, value: 40}
	b5 := &BST{key: 5, value: 50}
	b6 := &BST{key: 6, value: 60}
	b7 := &BST{key: 7, value: 70}
	b8 := &BST{key: 8, value: 80}
	b9 := &BST{key: 9, value: 90}

	b1.setRight(b2)
	b2.setRight(b3)
	b3.setRight(b4)
	b4.setRight(b5)
	b5.setRight(b6)
	b6.setRight(b7)
	b7.setRight(b8)
	b8.setRight(b9)

	return b1
}

func createBST3() *BST {
	b1 := &BST{key: 1, value: 10}
	b2 := &BST{key: 2, value: 20}
	b3 := &BST{key: 3, value: 30}
	b4 := &BST{key: 4, value: 40}
	b5 := &BST{key: 5, value: 50}
	b6 := &BST{key: 6, value: 60}
	b7 := &BST{key: 7, value: 70}
	b8 := &BST{key: 8, value: 80}
	b9 := &BST{key: 9, value: 90}

	b5.setLeft(b3)
	b5.setRight(b7)
	b3.setLeft(b2)
	b3.setRight(b4)
	b2.setLeft(b1)
	b7.setLeft(b6)
	b7.setRight(b8)
	b8.setRight(b9)

	return b5
}

func createBST4() *BST {
	b := &BST{key: 1, value: 10}
	last := b

	for i := 2; i < 100000; i++ {
		n := &BST{key: i, value: i * 10}
		last.setRight(n)

		last = n
	}

	return b
}

func createBST5() *BST {
	b := &BST{key: 99999, value: 999990}
	last := b

	for i := 99998; i > 0; i-- {
		n := &BST{key: i, value: i * 10}
		last.setLeft(n)

		last = n
	}

	return b
}

func createBST6() *BST {
	b := &BST{key: 50000, value: 500000}
	last := b

	for i := 49999; i > 0; i-- {
		n := &BST{key: i, value: i * 10}
		last.setLeft(n)

		last = n
	}

	last = b

	for i := 50001; i < 100000; i++ {
		n := &BST{key: i, value: i * 10}
		last.setRight(n)

		last = n
	}

	return b
}

func TestK(t *testing.T) {
	bst1 := createBST1()
	bst2 := createBST2()
	bst3 := createBST3()
	bst4 := createBST4()

	check(t, bst1, 1, 10)
	check(t, bst1, 2, 20)
	check(t, bst1, 3, 30)
	check(t, bst1, 4, 40)
	check(t, bst1, 5, 50)
	check(t, bst1, 6, 60)
	check(t, bst1, 7, 70)
	check(t, bst1, 8, 80)
	check(t, bst1, 9, 90)

	check(t, bst2, 1, 10)
	check(t, bst2, 2, 20)
	check(t, bst2, 3, 30)
	check(t, bst2, 4, 40)
	check(t, bst2, 5, 50)
	check(t, bst2, 6, 60)
	check(t, bst2, 7, 70)
	check(t, bst2, 8, 80)
	check(t, bst2, 9, 90)

	check(t, bst3, 1, 10)
	check(t, bst3, 2, 20)
	check(t, bst3, 3, 30)
	check(t, bst3, 4, 40)
	check(t, bst3, 5, 50)
	check(t, bst3, 6, 60)
	check(t, bst3, 7, 70)
	check(t, bst3, 8, 80)
	check(t, bst3, 9, 90)

	check(t, bst4, 1, 10)
	check(t, bst4, 2, 20)
	check(t, bst4, 3, 30)
	check(t, bst4, 4, 40)
	check(t, bst4, 5, 50)
	check(t, bst4, 6, 60)
	check(t, bst4, 7, 70)
	check(t, bst4, 8, 80)
	check(t, bst4, 9, 90)

	check(t, bst4, 99991, 999910)
	check(t, bst4, 99992, 999920)
	check(t, bst4, 99993, 999930)
	check(t, bst4, 99994, 999940)
	check(t, bst4, 99995, 999950)
	check(t, bst4, 99996, 999960)
	check(t, bst4, 99997, 999970)
	check(t, bst4, 99998, 999980)
	check(t, bst4, 99999, 999990)
}

func check(t *testing.T, b *BST, k, value int) {
	t.Helper()

	checkCall(t, b, k, value, count)
	// 此處可添加其他解法的實現
}

func checkCall(t *testing.T, b *BST, k, value int, find func(bst *BST, kth int) int) {
	t.Helper()

	got := find(b, k)
	if got != value {
		t.Fatalf("want:%d, got:%d", value, got)
	}
}
發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章