Amazon Online Assessment (OA) - Distance Between Nodes in BST

Given a list of unique integers nums.

Construct a BST from it (you need to insert nodes one-by-one with the given order to get the BST)

Find the distance between two nodes node1 and node2.

Distance is the number of edges between two nodes.

If any of the given nodes does not appear in the BST, return -1.


The input consists of three arguments:

nums: a list of integers representing the nodes list

node1: a integer representing the nodes

node2: a integer representing the nodes


return a integer representing distance between two nodes node1 and node2


Example 1:


nums = [2, 1, 3]

node1 = 1

node2 = 3

Output: 2


Try it yourself


from __future__ import annotations
from dataclasses import dataclass
from math import inf
from random import choices, shuffle
from typing import List, Optional

class Node:
    value: int
    left: Tree = None
    right: Tree = None

Tree = Optional[Node]

def insert(tree: Tree, val: int) -> Tree:
    if tree is None:
        tree = Node(val)
    elif val < tree.value:
        tree.left = insert(tree.left, val)
        tree.right = insert(tree.right, val)
    return tree

def find(tree: Tree, val: int) -> Optional[List[int]]:
    path = []
    while tree:
        if val == tree.value:
            return path
        if val < tree.value:
            tree = tree.left
            tree = tree.right
    return None

def bst_distance(nums: List[int], a: int, b: int) -> int:
    >>> bst_distance([2, 1, 3], 1, 3)
    >>> bst_distance([2, 1, 3], 1, 4)
    >>> bst_distance([2, 1, 3], 2, 1)
    >>> bst_distance([2, 1, 3], 3, 3)
    tree = None
    for n in nums:
        tree = insert(tree, n)
    left = find(tree, a)
    right = find(tree, b)
    if left is None or right is None:
        return -1
    return len(set(left).symmetric_difference(right))