This is purely a graph problem. We are given a binary tree (via a root node), a target node, and an integer K. We need to return a list of nodes at K distance from the target node. For the below example, the target node (marked yellow) is 5 and the answer will be [7, 4, 1] (marked blue), because K = 2.
Below are the cases to handle for solving problem:
The targets children at distance K.
The targets parent node being at distance K.
Other nodes being l+m distant from target: where l is the distance from target to a parent node par, and m is the distance from par to the node. (See the picture above)
If K is even, then you can say that target node is at distance K from itself. We are just assuming that we will always exclude the target node.
Step 2: Build a graph
It can be seen that, there is no escape from traversing from the target node to it’s parent for solving this problem. But it cannot be done in a tree data structure.
So the first step is to convert this tree to a bidirectional unweighted graph. We can traverse the tree and every time we go to a child node, we will assign the parent to the child. This can be done both recursively or iteratively.
Unlike other statically typed programming languages, python has the ability to create class attributes outside of class. We will use this ability for convenience.
1
2
3
4
5
6
7
8
# Time Complexity: O(n)defdfs(node,parent=None):ifnode:# Check if the node is not Nonenode.parent=parent# Create a parent attributedfs(node.left,node)dfs(node.right,node)dfs(root)# Now all nodes in the graph has parent attribute
Step 3: BFS
Now that we have the graph, we can easily apply a BFS (Breadth First Search) from the target to all nodes in the graph, check for the distance and if it is K, log it in a list. Finally, return the list. Again, we can keep the distance saved within the node itself.
If we find a node at distance K, BFS property tells that we have all the nodes at distance K already in the queue. So we can stop BFS and return the queue
# Time Complexity: O(n*n)target.dist=0# distance from itself is zeroqueue=[target]# BFS queueseen={target}# A set to check if visitedwhilequeue:# run until queue is emptyifqueue[0].dist==K:# Found one of our answer nodes# Return the queueans=[]fornodeinqueue:ans.append(node.val)returnansnode=queue.pop(0)# Takes O(len(queue))# Go through the neighbors of current nodeforneighborin(node.left,node.right,node.parent):ifneighborandneighbornotinseen:neighbor.dist=node.dist+1seen.add(neighbor)queue.append(neighbor)
Step 4: Handle corner cases
In this problem, there can be 2 corner cases.
There will be no node at that distance (the graph is very small compared to K).
The graph is empty.
In both cases, we will need to return an empty array.
1
return[]
Step 4: Optimize
The complete algorithm currently runs in O(n + n*n) = O(n*n) complexity. Notice that, we are using a list to do a queue’s job in BFS. Python has a built in queue and deque module which have constant access time for all operations. We can use deque for this purpose. Also in some places, we can use list comprehensions which will not effect on complexity, but may have some interpreter optimizations.
Finally, the time complexity will be O(n) and space complexity is also O(n) where n is the number of nodes in the tree. This is the best because all the n nodes must be visited once to solve the problem.
importcollections# Definition for a binary tree node.classTreeNode:def__init__(self,x):self.val=xself.left=Noneself.right=NoneclassSolution(object):defdistanceK(self,root:TreeNode,target:TreeNode,K:int)->[int]:defdfs(node,parent=None):ifnode:# Check if the node is not Nonenode.parent=parent# Create a parent attributedfs(node.left,node)dfs(node.right,node)dfs(root)# Now all nodes in the graph has parent attributetarget.dist=0# distance from itself is zeroqueue=collections.deque([target])# BFS queueseen={target}# A set to check if visitedwhilequeue:# run until queue is emptyifqueue[0].dist==K:# Found one of our answer nodesreturn[node.valfornodeinqueue]node=queue.popleft()# Go through the neighbors of current nodeforneighborin(node.left,node.right,node.parent):ifneighborandneighbornotinseen:neighbor.dist=node.dist+1seen.add(neighbor)queue.append(neighbor)return[]
/**
* Definition for a binary tree node.
* struct TreeNode {
* int val;
* TreeNode *left;
* TreeNode *right;
* TreeNode(int x) : val(x), left(NULL), right(NULL) {}
* };
*/classSolution{public:vector<int>distanceK(TreeNode*root,TreeNode*target,intK){map<TreeNode*,TreeNode*>parents;// Initial parents for all nodes
dfs(parents,root,nullptr);queue<pair<TreeNode*,int>>q;set<TreeNode*>seen;q.push({target,0});// Distance from itself is zero
seen.insert(target);// To check if visited
while(!q.empty()){// run until queue is empty
if(q.front().second==K){vector<int>result;while(!q.empty()){result.push_back(q.front().first->val);q.pop();}returnresult;}// Go through the neighbors of current node
autocur=q.front();q.pop();autonode=cur.first;autodist=cur.second;TreeNode*neighbors[]={node->left,node->right,parents[node]};for(autoneighbor:neighbors){if(neighbor!=nullptr&&seen.find(neighbor)==seen.end()){autonew_dist=dist+1;seen.insert(neighbor);q.push({neighbor,new_dist});}}}return{};}private:voiddfs(map<TreeNode*,TreeNode*>&parents,TreeNode*node,TreeNode*parent){if(node!=nullptr){// Check if the node is not Node
parents[node]=parent;// Create a parent
dfs(parents,node->left,node);dfs(parents,node->right,node);}}};
importjava.util.*;importjava.lang.*;/**
* Definition for a binary tree node.
* public class TreeNode {
* int val;
* TreeNode left;
* TreeNode right;
* TreeNode(int x) { val = x; }
* }
*/classPair{publicTreeNodefirst;publicIntegersecond;publicPair(TreeNodefirst,Integersecond){this.first=first;this.second=second;}}classSolution{publicList<Integer>distanceK(TreeNoderoot,TreeNodetarget,intK){Map<TreeNode,TreeNode>parents=newHashMap<>();dfs(parents,root,null);// Initial parents for all nodesQueue<Pair>queue=newLinkedList<>();Set<TreeNode>seen=newHashSet<>();queue.add(newPair(target,0));// Distance from itself is zeroseen.add(target);// To check if visitedwhile(!queue.isEmpty()){// run until queue is emptyif(queue.peek().second==K){// Found one of our answer nodesList<Integer>result=newArrayList<>();while(!queue.isEmpty())result.add(queue.remove().first.val);returnresult;}// Go through the neighbors of current nodePaircur=queue.remove();TreeNodenode=cur.first;intdist=cur.second;TreeNode[]neighbors={node.left,node.right,parents.get(node)};for(TreeNodeneighbor:neighbors){if(neighbor!=null&&!seen.contains(neighbor)){intnew_dist=dist+1;seen.add(neighbor);queue.add(newPair(neighbor,new_dist));}}}returnnewArrayList<Integer>();}privatevoiddfs(Map<TreeNode,TreeNode>parents,TreeNodenode,TreeNodeparent){if(node!=null){// Check if the node is not Nodeparents.put(node,parent);// Create a parent attributedfs(parents,node.left,node);dfs(parents,node.right,node);}}}