distanceK.swift

func distanceK(_ root: TreeNode?, _ target: TreeNode?, _ K: Int) -> [Int] {
        guard let root = root else { return [] }
        var parentDict = [TreeNode: TreeNode]()
        dfsFillParentDict(root, parentDict)
        
        var queue = [root, nil]
        var level = 0
        var visited = Set<TreeNode>()
        while level != K && !queue.isEmpty {
            if let node = queue.removeFirst() {
                if let left = node.left {
                    queue.append(left)
                }
                if let right = node.right {
                    queue.append(right)
                }
                if let parent = parentDict[node] {
                    queue.append(parent)
                }
                visited.insert(node)
            } else {
                if !queue.isEmpty{
                    queue.append(nil)
                    level += 1
                }
            }
        }
        return queue
    }
    
    func dfsFillParentDict(_ root: TreeNode?, _ parentDict: inout [TreeNode: TreeNode]) {
        guard let root = root else { return }
        dfsFillParentDict(root.left)
        dfsFillParentDict(root.right)
        if let left = root.left {
            parentDict[left] = root
        }
        if let right = root.right {
            parentDict[right] = root
        }
        
    }