616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121
from collections import deque

# Class representing a binary tree node
class TreeNode:
    # Constructor to initialize the node
    def __init__(self, x):
        # Value of the node
        self.val = x
        # Left child of the node
        self.left = None
        # Right child of the node
        self.right = None

class Solution:
    # Function to mark parents of nodes in the tree
    def markParents(self, root, parentTrack):
        # Use a queue for BFS traversal
        q = deque([root])

        # Continue until the queue is empty
        while q:
            # Remove the front node
            current = q.popleft()

            # If the left child exists
            if current.left:
                # Record parent of left child
                parentTrack[current.left] = current
                # Add left child to queue
                q.append(current.left)

            # If the right child exists
            if current.right:
                # Record parent of right child
                parentTrack[current.right] = current
                # Add right child to queue
                q.append(current.right)

    # Function to find all nodes at distance k from target
    def distanceK(self, root, target, k):
        # Dictionary to store parent references
        parentTrack = {}
        # Fill parent references
        self.markParents(root, parentTrack)

        # Set to track visited nodes
        visited = set()
        # Queue for BFS starting at target node
        q = deque([target])
        # Mark target as visited
        visited.add(target)

        # Variable to track current BFS level
        currLevel = 0

        # Continue BFS until reaching distance k
        while q:
            # Number of nodes in current level
            size = len(q)

            # If current level is equal to k, stop
            if currLevel == k:
                break
            currLevel += 1

            # Process all nodes in this level
            for _ in range(size):
                # Remove node from queue
                current = q.popleft()

                # Check left child
                if current.left and current.left not in visited:
                    # Add left child
                    q.append(current.left)
                    # Mark as visited
                    visited.add(current.left)

                # Check right child
                if current.right and current.right not in visited:
                    # Add right child
                    q.append(current.right)
                    # Mark as visited
                    visited.add(current.right)

                # Check parent
                if current in parentTrack and parentTrack[current] not in visited:
                    # Add parent
                    q.append(parentTrack[current])
                    # Mark as visited
                    visited.add(parentTrack[current])

        # Collect all nodes at distance k
        return [node.val for node in q]

# Driver code
if __name__ == "__main__":
    # Build the binary tree
    root = Tr