LeetCode面试题 04.06. 后继者

LeetCode面试题 04.06. 后继者

题目描述
设计一个算法，找出二叉搜索树中指定节点的“下一个”节点（也即中序后继）。

如果指定节点没有对应的“下一个”节点，则返回null。

示例 1:

输入: root = [2,1,3], p = 1

  2
 / \
1   3

输出: 2

示例 2:

输入: root = [5,3,6,2,4,null,null,1], p = 6

      5
     / \
    3   6
   / \
  2   4
 /   
1

输出: null

解题思路

1. 由于题目说了是二叉搜索树
2. 因此直接将树遍历一遍, 将每个节点的值用数组存起来
3. 排序之后, 找到指定节点
4. 然后dfs遍历找到指定节点

AC代码

class Solution {
public:
	TreeNode* ans=NULL;
	void dfs(TreeNode* root, int n) {
		if (root && !ans) {
			if (root->val == n) ans = root;
			else if (root->val > n) dfs(root->left, n);
			else dfs(root->right, n);
		}
	}
	TreeNode* inorderSuccessor(TreeNode* root, TreeNode* p) {
		queue<TreeNode*>q;
		vector<int> s;
		if (root) q.push(root);
		while (!q.empty()) {
			auto t = q.front();
			q.pop();
			s.push_back(t->val);
			if (t->left) q.push(t->left);
			if (t->right) q.push(t->right);
		}
		sort(s.begin(), s.end());
		auto it = upper_bound(s.begin(), s.end(), p->val);
		if (it == s.end())return NULL;
		dfs(root,*it);
		return ans;
	}
};