# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# +
# Assignment - Find All Nodes Distance K in Binary Tree
# Time: O(n) -> Two-pass approach
# Space: O(n)

from collections import deque

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

def print_k_leaf(root, k):
    def dfs(root, parent = None):
        if not root:    return root
        root.parent = parent
        if not root.left and not root.right:
            queue.append(root)
        dfs(root.left, root)
        dfs(root.right, root)
    
    queue = []
    dfs(root)
    res = []
    seen = set()
    for node in queue:
        K = k
        while K and node:
            node = node.parent
            K -= 1
        if not K and node and node not in seen:
            res.append(node.val)
            seen.add(node)
    return res

if __name__=='__main__':
    root = Node(3)
    root.left = Node(8)
    root.right = Node(9)
    root.left.left = Node(11)
    root.left.right = Node(7)
    root.left.right.left = Node(6)
    root.left.right.right = Node(12)
    root.right.left = Node(8)
    root.right.right = Node(3)

    print(print_k_leaf(root, 2))


# +
# Time: O(n) -> One-pass approach
# Space: O(n)

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

def print_k_leaf(node, k, res):
    def dfs(root, i = -1, arr = [], seen = []):
        if not root:    return root
        arr.append(root.val)
        seen.append(False)
        i += 1
        if not root.left and not root.right and i >= k:
            if not seen[i - k]:
                res.append(arr[i - k])
                seen[i - k] = True
        dfs(root.left, i, arr, seen)
        dfs(root.right, i, arr, seen)
        arr.pop()
        seen.pop()
    dfs(root)

if __name__=='__main__':
    root = Node(3)
    root.left = Node(8)
    root.right = Node(9)
    root.left.left = Node(11)
    root.left.right = Node(7)
    root.left.right.left = Node(6)
    root.left.right.right = Node(12)
    root.right.left = Node(8)
    root.right.right = Node(3)
    
    res = []
    print_k_leaf(root, 2, res)
    print(res)
# -



StarcoderdataJupyterScriptsDedupFiltered