{}
//Variant of all nodes at k distance
/**
* Variant: was: you’re given the root node of a binary tree, the value N of a target node, a distance K and a target sum T. Find all sets of nodes at distance K from node N which sum to T.
*
*/
public class TreeNode {
public var val: Int
public var left: TreeNode?
public var right: TreeNode?
init(_ val: Int, _ left: TreeNode? = nil, _ right: TreeNode? = nil) {
self.val = val
self.left = left
self.right = right
}
}
extension TreeNode: Hashable {
public static func == (lhs: TreeNode, rhs: TreeNode) -> Bool {
return lhs === rhs
}
public func hash(into hasher: inout Hasher) {
hasher.combine(ObjectIdentifier(self))
}
}
class Solution {
//we need to look at the tree as a graph instead of tree
var graph = [TreeNode: [TreeNode]]() //adj list
func distanceKEqualSum(_ root: TreeNode?, _ targetVal: Int, _ k: Int, _ targetSum: Int) -> [[Int]] {
guard let root = root, let target = findTarget(root, targetVal) else { return [] }
buildGraph(root, nil)
var result = [[Int]]()
var nodes = [Int]() //nodes at K distance
var visited = Set<TreeNode>()
visited.insert(target)
distanceKHelper(target, 0, k, &visited, &nodes)
print(nodes)
//we need to perform subset sum in the nodes list to get those nodes that sum to targetSum
subsetSumHelper(nodes, 0, targetSum, [], &result)
return result
}
private func findTarget(_ root: TreeNode?, _ targetVal: Int) -> TreeNode? {
guard let root = root else { return nil }
if root.val == targetVal { return root }
return findTarget(root.left, targetVal) ?? findTarget(root.right, targetVal)
}
private func subsetSumHelper(_ nums: [Int], _ index: Int, _ targetSum: Int, _ path: [Int], _ result: inout [[Int]]) {
if targetSum == 0 {
result.append(path)
return
}
if targetSum < 0 || index >= nums.count { return }
//include element at index
subsetSumHelper(nums, index + 1, targetSum - nums[index], path + [nums[index]], &result)
//exclude element at index
subsetSumHelper(nums, index + 1, targetSum, path, &result)
}
private func distanceKHelper(_ node: TreeNode, _ distance: Int, _ k: Int, _ visited: inout Set<TreeNode>, _ result: inout [Int]) {
if distance == k {
result.append(node.val)
return
}
for neighbor in graph[node, default: []] {
guard visited.contains(neighbor) == false else { continue }
visited.insert(neighbor)
distanceKHelper(neighbor, distance + 1, k, &visited, &result)
}
}
private func buildGraph(_ root: TreeNode?, _ parent: TreeNode?) {
if let root = root, let parent = parent {
graph[root, default: []].append(parent)
graph[parent, default: []].append(root)
}
if let left = root?.left {
buildGraph(left, root)
}
if let right = root?.right {
buildGraph(right, root)
}
}
}
let root = TreeNode(5,
TreeNode(3,
TreeNode(1),
TreeNode(4)
),
TreeNode(8,
TreeNode(7),
TreeNode(9)
)
)
let solution = Solution()
let results = solution.distanceKEqualSum(root, 3, 1, 5)
print(results) // Output: possible [[1, 9], [4, 6]] etc., depending on tree and target