基于TreeMap与DFS实现二叉树顶视图遇错求助

我尝试用TreeMap结合递归DFS实现二叉树的顶视图功能，思路是让TreeMap以列号为键，通过putIfAbsent方法只存储每列最上层的节点值。但所有测试用例都未通过，目前有几个疑问：

• 递归方法里的row参数没用到，能不能直接去掉？
• 我以为DFS的遍历顺序会让同一列的上层节点先被访问，所以putIfAbsent能保证存的是最上层节点，但实际结果错误，想知道问题出在哪。

我的实现代码 #

递归遍历方法：

public static void verticalNodes(TreeNode root, TreeMap<Integer, Integer> solve, int row, int col) {
    if (root == null) {
        return;
    }
    solve.putIfAbsent(col, root.data);
    verticalNodes(root.left, solve, row + 1, col - 1);
    verticalNodes(root.right, solve, row + 1, col + 1);        
}

生成顶视图列表的方法：

public static List<Integer> getTopView(TreeNode root) {
    List<Integer> ans = new ArrayList<>();
    if (root == null) {
        return ans;
    }
    TreeMap<Integer, Integer> solve = new TreeMap<>();
    verticalNodes(root, solve, 0, 0);
    for (int val : solve.values()) {
        ans.add(val);
    }        
    return ans;
}

问题根源 #

DFS的遍历顺序无法保证同一列中最上层（行号最小）的节点被优先访问。举个例子：假设某棵树的右子树左分支存在一个列号为-1的节点，且这个节点会被DFS先遍历到；而同一列号-1的上层节点在左子树的右分支，会被后遍历到。此时putIfAbsent已经把深层节点的值存进了TreeMap，后续上层节点无法覆盖，导致顶视图结果错误。

你的row参数确实不能去掉，反而需要用它来判断当前节点是否是该列的最上层节点——只有当当前节点的行号小于该列已记录的最小行号时，才更新TreeMap的值。

修正方案 #

方案1：改进DFS，维护每列的最小行号

新增一个TreeMap来记录每个列对应的最小行号，递归时比较当前行号与已记录的行号，只有当前行号更小时才更新节点值：

public static void verticalNodes(TreeNode root, TreeMap<Integer, Integer> solve, TreeMap<Integer, Integer> minRowMap, int row, int col) {
    if (root == null) {
        return;
    }
    // 若列未记录，或当前行号比已记录的最小行号更小，则更新
    if (!minRowMap.containsKey(col) || row < minRowMap.get(col)) {
        solve.put(col, root.data);
        minRowMap.put(col, row);
    }
    verticalNodes(root.left, solve, minRowMap, row + 1, col - 1);
    verticalNodes(root.right, solve, minRowMap, row + 1, col + 1);        
}

public static List<Integer> getTopView(TreeNode root) {
    List<Integer> ans = new ArrayList<>();
    if (root == null) {
        return ans;
    }
    TreeMap<Integer, Integer> solve = new TreeMap<>();
    TreeMap<Integer, Integer> minRowMap = new TreeMap<>();
    verticalNodes(root, solve, minRowMap, 0, 0);
    for (int val : solve.values()) {
        ans.add(val);
    }        
    return ans;
}

方案2：改用BFS（层序遍历）更高效

BFS是按层遍历，同一列的上层节点会被优先访问，天然符合顶视图的需求，不需要维护行号判断：

import java.util.Queue;
import java.util.LinkedList;

public static List<Integer> getTopView(TreeNode root) {
    List<Integer> ans = new ArrayList<>();
    if (root == null) {
        return ans;
    }
    TreeMap<Integer, Integer> solve = new TreeMap<>();
    Queue<Pair<TreeNode, Integer>> queue = new LinkedList<>();
    queue.add(new Pair<>(root, 0));
    
    while (!queue.isEmpty()) {
        Pair<TreeNode, Integer> curr = queue.poll();
        TreeNode node = curr.getKey();
        int col = curr.getValue();
        
        // 只在列未记录时存入，保证是最上层节点
        solve.putIfAbsent(col, node.data);
        
        if (node.left != null) {
            queue.add(new Pair<>(node.left, col - 1));
        }
        if (node.right != null) {
            queue.add(new Pair<>(node.right, col + 1));
        }
    }
    
    ans.addAll(solve.values());
    return ans;
}

// 若Java版本无内置Pair，可自定义：
class Pair<K, V> {
    private K key;
    private V value;
    public Pair(K key, V value) {
        this.key = key;
        this.value = value;
    }
    public K getKey() { return key; }
    public V getValue() { return value; }
}

内容的提问来源于stack exchange，提问作者Garvit