# mohitathwani/SwiftCodingChallenges

Fetching contributors…
Cannot retrieve contributors at this time
147 lines (120 sloc) 3.53 KB
 //https://leetcode.com/problems/all-nodes-distance-k-in-binary-tree/solution/ //We are given a binary tree (with root node root), a target node, and an integer value `K`. // //Return a list of the values of all nodes that have a distance K from the target node. The answer can be returned in any order. // // // //Example 1: // //Input: root = [3,5,1,6,2,0,8,null,null,7,4], target = 5, K = 2 //Output: [7,4,1] //Explanation: //The nodes that are a distance 2 from the target node (with value 5) //have values 7, 4, and 1. // //Note that the inputs "root" and "target" are actually TreeNodes. //The descriptions of the inputs above are just serializations of these objects. // //Note: // //The given tree is non-empty. //Each node in the tree has unique values 0 <= node.val <= 500. //The target node is a node in the tree. //0 <= K <= 1000. public class TreeNode { public var val: Int public var left: TreeNode? public var right: TreeNode? public init(_ val: Int) { self.val = val self.left = nil self.right = nil } } extension TreeNode: Hashable, CustomStringConvertible { public var description: String { return "\(val)" } public var hashValue: Int { return val } public static func == (lhs: TreeNode, rhs: TreeNode) -> Bool { return lhs.val == rhs.val } } func distanceK(_ root: TreeNode?, _ target: TreeNode?, _ K: Int) -> [Int] { guard let root = root, let target = target else { return [Int]() } var parentMap = [TreeNode:TreeNode?]() var result = [Int]() func dfs(currentNode: TreeNode?, parentNode: TreeNode?) { guard let currentNode = currentNode else { return } parentMap[currentNode] = parentNode dfs(currentNode: currentNode.left, parentNode: currentNode) dfs(currentNode: currentNode.right, parentNode: currentNode) } dfs(currentNode: root, parentNode: nil) func bfs(startNode: TreeNode) { var q = [TreeNode]() var seenSet = [TreeNode]() q.append(startNode) seenSet.append(startNode) var distanceMap = [TreeNode:Int]() distanceMap[startNode] = 0 if K == 0 { result.append(startNode.val) return } while !q.isEmpty { let node = q.removeFirst() if let left = node.left, !seenSet.contains(left) { q.append(left) seenSet.append(left) distanceMap[left] = distanceMap[node]! + 1 if distanceMap[left]! == K { result.append(left.val) } } if let right = node.right, !seenSet.contains(right) { q.append(right) seenSet.append(right) distanceMap[right] = distanceMap[node]! + 1 if distanceMap[right]! == K { result.append(right.val) } } if let optionalParent = parentMap[node], let parent = optionalParent, !seenSet.contains(parent) { q.append(parent) seenSet.append(parent) distanceMap[parent] = distanceMap[node]! + 1 if distanceMap[parent]! == K { result.append(parent.val) } } } } bfs(startNode: target) return result } let root = TreeNode(3) let fiveNode = TreeNode(5) let oneNode = TreeNode(1) let sixNode = TreeNode(6) let twoNode = TreeNode(2) let zeroNode = TreeNode(0) let eightNode = TreeNode(8) let sevenNode = TreeNode(7) let fourNode = TreeNode(4) root.left = fiveNode root.right = oneNode fiveNode.left = sixNode fiveNode.right = twoNode oneNode.left = zeroNode oneNode.right = eightNode twoNode.left = sevenNode twoNode.right = fourNode distanceK(root, root, 0)