如何让递归实现的堆add函数返回新增节点的深度？

问题描述 #

用带节点的二叉树实现了一个堆，现有add函数能将新节点插入堆的正确位置，但无法让函数返回新增元素最终所在节点的深度（即节点在树中的层级）。核心问题是函数采用递归实现，深度计数器难以管理，且因前期设计问题（通过交换值而非创建新节点），无法为Node类添加int类型的深度属性。

示例输出 #

System.out.println(tree.add(5)); //输出0
System.out.println(tree.add(2)); //输出0
System.out.println(tree.add(7)); //输出0
System.out.println(tree.add(1)); //输出0
System.out.println(tree.add(100)); //输出2

现有代码 #

public class HeapBinaryTree {
    Node root;

    public class Node {
        public Integer key;
        public Node left, right;
        public int size;
        public int value;

        public Node(Integer key) {
            this.key = key;
            this.value = value;
            this.left = this.right = null;
        }
        private Integer add(Integer key) { 
            Node cur = this;
            cur.size++;

            if (key < cur.key) {
                int temp = cur.key;
                cur.key = key;
                key = temp;
            }
            if (cur.left == null) {
                cur.left = new Node(key);
                return;
            }
            if (cur.right == null) {
                cur.right = new Node(key);
                return;
            }
            if (cur.left.size < cur.right.size) {
                cur.left.add(key);
            } else {
                cur.right.add(key);
            }
            return null;
        }
    }

    public Integer addFirst(Integer key) { //Node类外部的方法
        if (root == null) {
            root = new Node(key);
            return 0;
        } else {
            int ret = root.add(key);
            return ret; //我希望此处返回深度
        }
    }
}

解决方案 #

核心思路 #

递归调用时传递当前节点的深度，同时区分两种场景处理返回值：

1. 发生值交换：新元素被上浮到当前节点，此时当前节点的深度就是该元素最终的深度，递归处理被替换的旧值即可，返回当前节点深度。
2. 直接插入子节点：新元素无需上浮，插入到空的左/右节点时，返回当前节点深度+1（子节点的深度）；若子节点都存在，递归插入到子树并返回子树返回的深度。

修改后的代码 #

1. 调整Node类的add方法，添加深度参数

private Integer add(Integer key, int currentDepth) { 
    Node cur = this;
    cur.size++;

    // 堆调整：若新元素比当前节点小，交换值并返回当前节点深度
    if (key < cur.key) {
        int temp = cur.key;
        cur.key = key;
        // 递归插入被替换的旧值，不影响返回值
        cur.add(temp, currentDepth + 1);
        return currentDepth;
    }

    // 插入左节点，返回子节点深度
    if (cur.left == null) {
        cur.left = new Node(key);
        return currentDepth + 1;
    }
    // 插入右节点，返回子节点深度
    if (cur.right == null) {
        cur.right = new Node(key);
        return currentDepth + 1;
    }

    // 递归插入到子树，返回子树的插入深度
    if (cur.left.size < cur.right.size) {
        return cur.left.add(key, currentDepth + 1);
    } else {
        return cur.right.add(key, currentDepth + 1);
    }
}

2. 调整外部addFirst方法，传递初始深度

public Integer addFirst(Integer key) { 
    if (root == null) {
        root = new Node(key);
        return 0;
    } else {
        // 根节点深度为0，传入递归函数
        return root.add(key, 0);
    }
}

匹配示例输出的调整说明 #

若要完全匹配你给出的示例输出，说明你的堆是大顶堆（新元素大于当前节点时上浮），只需将交换条件改为key > cur.key即可：

if (key > cur.key) {
    int temp = cur.key;
    cur.key = key;
    cur.add(temp, currentDepth + 1);
    return currentDepth;
}

调整后执行测试用例将完全符合示例输出。

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