Nodes at given distance in binary tree (Amazon SDE-2)

Question

Given a binary tree, a target node in the binary tree, and an integer value k, find all the nodes that are at distance k from the given target node. No parent pointers are available.

link to the problem on GFG: LINK

Example 1:

Input :      
          20
        /    \
      8       22 
    /   \
   4    12 
       /   \
      10    14

Target Node = 8
K = 2

Output: 10 14 22

Explanation: The three nodes at distance 2
from node 8 are 10, 14, 22.

My code

from collections import defaultdict
class solver:
    
    def __init__(self):
        self.vertList = defaultdict(list)
        
    def addEdge(self,u,v):
        self.vertList[u].append(v)
        
    def makeGraph(self,root):
        visited = set()
        queue = []
        queue.append(root)
        
        while len(queue) > 0:
            curr = queue.pop(0)
            visited.add(curr)
            
            if curr.left is not None and curr.left not in visited:
                self.vertList[curr.data].append(curr.left.data)
                self.vertList[curr.left.data].append(curr.data)

                queue.append(curr.left)
                
            if curr.right is not None and curr.right not in visited:
                self.vertList[curr.data].append(curr.right.data)
                self.vertList[curr.right.data].append(curr.data)

                queue.append(curr.right)

    def KDistanceNodes(self,root,target,k):
        self.makeGraph(root)
        dist = {}
        for v in self.vertList:
            dist[v] = 0

        visited2 = set()
        queue2 = []
        queue2.append(target)
        
        while len(queue2) > 0:
            curr = queue2.pop(0)
            visited2.add(curr)

            for nbr in self.vertList[curr]:
                if nbr not in visited2:
                    visited2.add(nbr)
                    queue2.append(nbr)
                    dist[nbr] = dist[curr] + 1
        ans = []
        for v in dist:
            if dist[v] == k:
                ans.append(str(v))
            
        return ans

#{ 
#  Driver Code Starts
#Initial Template for Python 3

from collections import deque

# Tree Node
class Node:
    def __init__(self, val):
        self.right = None
        self.data = val
        self.left = None

# Function to Build Tree
def buildTree(s):
    # Corner Case
    if (len(s) == 0 or s[0] == "N"):
        return None

    # Creating list of strings from input
    # string after spliting by space
    ip = list(map(str, s.split()))

    # Create the root of the tree
    root = Node(int(ip[0]))
    size = 0
    q = deque()

    # Push the root to the queue
    q.append(root)
    size = size + 1

    # Starting from the second element
    i = 1
    while (size > 0 and i < len(ip)):
        # Get and remove the front of the queue
        currNode = q[0]
        q.popleft()
        size = size - 1

        # Get the current node's value from the string
        currVal = ip[i]

        # If the left child is not null
        if (currVal != "N"):
            # Create the left child for the current node
            currNode.left = Node(int(currVal))

            # Push it to the queue
            q.append(currNode.left)
            size = size + 1
        # For the right child
        i = i + 1
        if (i >= len(ip)):
            break
        currVal = ip[i]

        # If the right child is not null
        if (currVal != "N"):
            # Create the right child for the current node
            currNode.right = Node(int(currVal))

            # Push it to the queue
            q.append(currNode.right)
            size = size + 1
        i = i + 1
    return root

if __name__ == "__main__":
    x = solver()
    t = int(input())
    for _ in range(t):
        line = input()
        target=int(input())
        k=int(input())

        root = buildTree(line)
        res = x.KDistanceNodes(root,target,k)
        
        for i in res:
            print(i, end=' ')
        print()

Input:

1 N 2 N 3 N 4 5
target = 5
k = 4

Its Correct output is:

1

And Your Code's output is:

[]

My logic:

-> First convert the tree into a undirected graph using BFS / level order traversal

-> Traverse the graph using BFS and calculate the distance

-> Return the nodes at k distance from the target

What I think: First of all in the given test case the tree representation seems to be in level order however, in the failing test case it doesn't look like level order or maybe my logic is wrong?

Input Format:

Custom input should have 3 lines. First line contains a string representing the tree as described below. Second line contains the data value of the target node. Third line contains the value of K.

The values in the string are in the order of level order traversal of the tree where, numbers denote node values, and a character “N” denotes NULL child.

For the above tree, the string will be: 1 2 3 N N 4 6 N 5 N N 7 N

Answer 1

The mistake is that the logic that is done in "init" should be done for every use-case (reinitializing self.vertList), not just once at the beginning.

Change:

def __init__(self):
    self.vertList = defaultdict(list)

to:

 def __init__(self):
    pass

and:

def KDistanceNodes(self,root,target,k):
    self.makeGraph(root)
    dist = {}
    ...

to:

def KDistanceNodes(self, root, target, k):
    self.vertList = defaultdict(list) # <-- move it here
    self.makeGraph(root)
    dist = {}
    ...

The original code kept accumulating the nodes and "remembered" the previous use-cases.

Also pay attention that you should return ans sorted, meaning that you should not append the numbers as strings so that you'll be able to sort them, change: ans.append(str(v)) to: ans.append(v) and return sorted(ans).