863. 二叉树中所有距离为 K 的结点

中等二叉树父指针表 + 广度优先搜索力扣原题 ↗

题目描述

输入二叉树的根节点 root、树中某个目标节点 target 和整数 k。树用层序数组表示，空位为 null；target 用节点值给出，树中节点值互不相同。

两个节点的距离是它们之间路径上的边数，路径既可以向下走到孩子，也可以向上走到父节点。

返回所有与 target 距离恰好为 k 的节点的值，顺序不限。没有这样的节点时返回空数组。

示例 1

输入：root = [3,5,1,6,2,0,8,null,null,7,4], target = 5, k = 2
输出：[7,4,1]

解释：节点 5 的两个孙子 7、4 距离为 2，向上经过根 3 到达的节点 1 距离也为 2。

示例 2

输入：root = [1], target = 1, k = 3
输出：[]

解释：树只有一个节点，不存在距离为 3 的节点。

数据范围

节点数在 [1, 500] 范围内。
0 <= Node.val <= 500
所有 Node.val 互不相同。
target 是树中的节点。
0 <= k <= 1000

解题思路：父指针表 + 广度优先搜索

树只能从父走向子，但距离可以往上走。先 DFS 一遍记下每个节点的父节点，树就变成了无向图：每个节点有左、右、父三个邻居。然后从 target 出发做 BFS，扩展 k 层，第 k 层的节点就是答案。用 seen 防止沿着刚走过的边走回去。

为什么正确

第 i 轮扩展后，frontier 恰好是与 target 距离为 i 的所有节点，seen 是距离 <= i 的所有节点。

复杂度分析

时间 O(n)，空间 O(n)（父指针表、seen 和队列）。

易错点

只向下搜索会漏掉经过父节点的路径（例如示例里的 1）；加了父指针后忘记 seen 又会在父子之间来回走，得到距离错误的节点。

参考代码

Python、Java、C++、Go、JavaScript、TypeScript 六种语言，逻辑一致。

class Solution:
    def distanceK(self, root: TreeNode, target: TreeNode, k: int) -> List[int]:
        parent = {}
        def dfs(node, par):
            if node:
                parent[node] = par
                dfs(node.left, node)
                dfs(node.right, node)
        dfs(root, None)
        seen = {target}
        frontier = [target]
        for _ in range(k):
            nxt = []
            for node in frontier:
                for nb in (node.left, node.right, parent[node]):
                    if nb and nb not in seen:
                        seen.add(nb)
                        nxt.append(nb)
            frontier = nxt
        return [node.val for node in frontier]

class Solution {
    Map<TreeNode, TreeNode> parent = new HashMap<>();

    void dfs(TreeNode node, TreeNode par) {
        if (node != null) {
            parent.put(node, par);
            dfs(node.left, node);
            dfs(node.right, node);
        }
    }

    public List<Integer> distanceK(TreeNode root, TreeNode target, int k) {
        parent.clear();
        dfs(root, null);
        Set<TreeNode> seen = new HashSet<>();
        seen.add(target);
        List<TreeNode> frontier = new ArrayList<>(List.of(target));
        for (int i = 0; i < k; i++) {
            List<TreeNode> nxt = new ArrayList<>();
            for (TreeNode node : frontier) {
                for (TreeNode nb : new TreeNode[]{node.left, node.right, parent.get(node)}) {
                    if (nb != null && !seen.contains(nb)) {
                        seen.add(nb);
                        nxt.add(nb);
                    }
                }
            }
            frontier = nxt;
        }
        List<Integer> res = new ArrayList<>();
        for (TreeNode node : frontier) res.add(node.val);
        return res;
    }
}

class Solution {
public:
    vector<int> distanceK(TreeNode* root, TreeNode* target, int k) {
        unordered_map<TreeNode*, TreeNode*> parent;
        function<void(TreeNode*, TreeNode*)> dfs = [&](TreeNode* node, TreeNode* par) {
            if (node) {
                parent[node] = par;
                dfs(node->left, node);
                dfs(node->right, node);
            }
        };
        dfs(root, nullptr);
        unordered_set<TreeNode*> seen{target};
        vector<TreeNode*> frontier{target};
        for (int i = 0; i < k; i++) {
            vector<TreeNode*> nxt;
            for (TreeNode* node : frontier) {
                for (TreeNode* nb : {node->left, node->right, parent[node]}) {
                    if (nb && !seen.count(nb)) {
                        seen.insert(nb);
                        nxt.push_back(nb);
                    }
                }
            }
            frontier = nxt;
        }
        vector<int> res;
        for (TreeNode* node : frontier) res.push_back(node->val);
        return res;
    }
};

func distanceK(root *TreeNode, target *TreeNode, k int) []int {
	parent := map[*TreeNode]*TreeNode{}
	var dfs func(node, par *TreeNode)
	dfs = func(node, par *TreeNode) {
		if node != nil {
			parent[node] = par
			dfs(node.Left, node)
			dfs(node.Right, node)
		}
	}
	dfs(root, nil)
	seen := map[*TreeNode]bool{target: true}
	frontier := []*TreeNode{target}
	for i := 0; i < k; i++ {
		nxt := []*TreeNode{}
		for _, node := range frontier {
			for _, nb := range []*TreeNode{node.Left, node.Right, parent[node]} {
				if nb != nil && !seen[nb] {
					seen[nb] = true
					nxt = append(nxt, nb)
				}
			}
		}
		frontier = nxt
	}
	res := []int{}
	for _, node := range frontier {
		res = append(res, node.Val)
	}
	return res
}

var distanceK = function(root, target, k) {
    const parent = new Map();
    const dfs = (node, par) => {
        if (node) {
            parent.set(node, par);
            dfs(node.left, node);
            dfs(node.right, node);
        }
    };
    dfs(root, null);
    const seen = new Set([target]);
    let frontier = [target];
    for (let i = 0; i < k; i++) {
        const nxt = [];
        for (const node of frontier) {
            for (const nb of [node.left, node.right, parent.get(node)]) {
                if (nb && !seen.has(nb)) {
                    seen.add(nb);
                    nxt.push(nb);
                }
            }
        }
        frontier = nxt;
    }
    return frontier.map(node => node.val);
};

function distanceK(root: TreeNode | null, target: TreeNode | null, k: number): number[] {
    const parent = new Map<TreeNode, TreeNode | null>();
    const dfs = (node: TreeNode | null, par: TreeNode | null): void => {
        if (node) {
            parent.set(node, par);
            dfs(node.left, node);
            dfs(node.right, node);
        }
    };
    dfs(root, null);
    const seen = new Set<TreeNode>([target!]);
    let frontier: TreeNode[] = [target!];
    for (let i = 0; i < k; i++) {
        const nxt: TreeNode[] = [];
        for (const node of frontier) {
            for (const nb of [node.left, node.right, parent.get(node)]) {
                if (nb && !seen.has(nb)) {
                    seen.add(nb);
                    nxt.push(nb);
                }
            }
        }
        frontier = nxt;
    }
    return frontier.map(node => node.val);
}

载入编辑器，自己跑一遍 →

其他解法

父指针法把树变成无向图，直接套广度优先模板；距离递归法利用树的结构在回溯时分类收集，省掉哈希表，但推理更细。

解法 2：递归返回到 target 的距离时间 O(n)，空间 O(h)

先写辅助函数 collect(node, d)，把 node 子树中深度恰好为 d 的节点加入答案。主递归 dfs(node) 返回 node 到 target 的距离（target 不在子树中返回 -1）：遇到 target 时调用 collect(target, k) 并返回 0；如果 target 在左子树且距离为 d，那么当 d = k 时当前节点就是答案，否则到右子树里收集深度为 k - d - 1 的节点；右子树同理。

不需要父指针表和 seen 集合，除答案外只用递归栈。但要把“向上走再拐到另一侧”的距离算清楚，分类讨论比把树当无向图做广度优先更容易出错。

class Solution:
    def distanceK(self, root: TreeNode, target: TreeNode, k: int) -> List[int]:
        ans = []

        def collect(node, d):
            if node is None or d < 0:
                return
            if d == 0:
                ans.append(node.val)
                return
            collect(node.left, d - 1)
            collect(node.right, d - 1)

        def dfs(node):
            if node is None:
                return -1
            if node is target:
                collect(node, k)
                return 0
            for child, other in ((node.left, node.right), (node.right, node.left)):
                d = dfs(child)
                if d != -1:
                    d += 1
                    if d == k:
                        ans.append(node.val)
                    else:
                        collect(other, k - d - 1)
                    return d
            return -1

        dfs(root)
        return ans

class Solution {
    List<Integer> ans;
    TreeNode target;
    int k;

    void collect(TreeNode node, int d) {
        if (node == null || d < 0) return;
        if (d == 0) {
            ans.add(node.val);
            return;
        }
        collect(node.left, d - 1);
        collect(node.right, d - 1);
    }

    int dfs(TreeNode node) {
        if (node == null) return -1;
        if (node == target) {
            collect(node, k);
            return 0;
        }
        int d = dfs(node.left);
        if (d != -1) {
            d++;
            if (d == k) ans.add(node.val);
            else collect(node.right, k - d - 1);
            return d;
        }
        d = dfs(node.right);
        if (d != -1) {
            d++;
            if (d == k) ans.add(node.val);
            else collect(node.left, k - d - 1);
            return d;
        }
        return -1;
    }

    public List<Integer> distanceK(TreeNode root, TreeNode target, int k) {
        ans = new ArrayList<>();
        this.target = target;
        this.k = k;
        dfs(root);
        return ans;
    }
}

class Solution {
public:
    vector<int> ans;
    TreeNode* target;
    int k;

    void collect(TreeNode* node, int d) {
        if (!node || d < 0) return;
        if (d == 0) {
            ans.push_back(node->val);
            return;
        }
        collect(node->left, d - 1);
        collect(node->right, d - 1);
    }

    int dfs(TreeNode* node) {
        if (!node) return -1;
        if (node == target) {
            collect(node, k);
            return 0;
        }
        int d = dfs(node->left);
        if (d != -1) {
            d++;
            if (d == k) ans.push_back(node->val);
            else collect(node->right, k - d - 1);
            return d;
        }
        d = dfs(node->right);
        if (d != -1) {
            d++;
            if (d == k) ans.push_back(node->val);
            else collect(node->left, k - d - 1);
            return d;
        }
        return -1;
    }

    vector<int> distanceK(TreeNode* root, TreeNode* target, int k) {
        ans.clear();
        this->target = target;
        this->k = k;
        dfs(root);
        return ans;
    }
};

func distanceK(root *TreeNode, target *TreeNode, k int) []int {
	ans := []int{}
	var collect func(node *TreeNode, d int)
	collect = func(node *TreeNode, d int) {
		if node == nil || d < 0 {
			return
		}
		if d == 0 {
			ans = append(ans, node.Val)
			return
		}
		collect(node.Left, d-1)
		collect(node.Right, d-1)
	}
	var dfs func(node *TreeNode) int
	dfs = func(node *TreeNode) int {
		if node == nil {
			return -1
		}
		if node == target {
			collect(node, k)
			return 0
		}
		for _, pair := range [2][2]*TreeNode{{node.Left, node.Right}, {node.Right, node.Left}} {
			d := dfs(pair[0])
			if d != -1 {
				d++
				if d == k {
					ans = append(ans, node.Val)
				} else {
					collect(pair[1], k-d-1)
				}
				return d
			}
		}
		return -1
	}
	dfs(root)
	return ans
}

var distanceK = function(root, target, k) {
    const ans = [];
    const collect = (node, d) => {
        if (node === null || d < 0) return;
        if (d === 0) {
            ans.push(node.val);
            return;
        }
        collect(node.left, d - 1);
        collect(node.right, d - 1);
    };
    const dfs = (node) => {
        if (node === null) return -1;
        if (node === target) {
            collect(node, k);
            return 0;
        }
        for (const [child, other] of [[node.left, node.right], [node.right, node.left]]) {
            let d = dfs(child);
            if (d !== -1) {
                d++;
                if (d === k) ans.push(node.val);
                else collect(other, k - d - 1);
                return d;
            }
        }
        return -1;
    };
    dfs(root);
    return ans;
};

function distanceK(root: TreeNode | null, target: TreeNode | null, k: number): number[] {
    const ans: number[] = [];
    const collect = (node: TreeNode | null, d: number): void => {
        if (node === null || d < 0) return;
        if (d === 0) {
            ans.push(node.val);
            return;
        }
        collect(node.left, d - 1);
        collect(node.right, d - 1);
    };
    const dfs = (node: TreeNode | null): number => {
        if (node === null) return -1;
        if (node === target) {
            collect(node, k);
            return 0;
        }
        for (const [child, other] of [[node.left, node.right], [node.right, node.left]]) {
            let d = dfs(child);
            if (d !== -1) {
                d++;
                if (d === k) ans.push(node.val);
                else collect(other, k - d - 1);
                return d;
            }
        }
        return -1;
    };
    dfs(root);
    return ans;
}

二叉树专题的其他题目

查看二叉树全部题目 →