面试题 04.03. 特定深度节点链表

面试题 04.03. 特定深度节点链表

/**
 * Definition for a binary tree node.
 * public class TreeNode {
 *     int val;
 *     TreeNode left;
 *     TreeNode right;
 *     TreeNode(int x) { val = x; }
 * }
 */
/**
 * Definition for singly-linked list.
 * public class ListNode {
 *     int val;
 *     ListNode next;
 *     ListNode(int x) { val = x; }
 * }
 */
class Solution {
    public ListNode[] listOfDepth(TreeNode tree) {
        if (tree == null)
            return new ListNode[0];
        List<ListNode> ans = new LinkedList<>();
        Queue<TreeNode> queue = new LinkedList<>();
        queue.add(tree);
        int currentLevel = 1;
        int nextLevel = 0;
        ListNode dummy = new ListNode();
        dummy.next = null;
        ListNode p = dummy;
        while (!queue.isEmpty()) {
            TreeNode currentNode = queue.remove();
            currentLevel--;
            // 尾插法：将当前层次的当前节点添加到指针p的后面，p随后后移
            ListNode node = new ListNode(currentNode.val);
            node.next = p.next;
            p.next = node;
            p = p.next;
            if (currentNode.left != null) {
                nextLevel++;
                queue.add(currentNode.left);
            }
            if (currentNode.right != null) {
                nextLevel++;
                queue.add(currentNode.right);
            }
            if (currentLevel == 0) {
                ans.add(dummy.next);
                currentLevel = nextLevel;
                nextLevel = 0;
                // 构造保存下一层的链表的头节点
                dummy = new ListNode();
                dummy.next = null;
                p = dummy;
            }
        } // while
        return ans.toArray(new ListNode[ans.size()]);
    }
}