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.

Input

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

Output

return a integer representing distance between two nodes node1 and node2

Examples

Example 1:

Input:

nums = [2, 1, 3]

node1 = 1

node2 = 3

Output: 2

Explanation:

Try it yourself

Solution

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


@dataclass
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)
    else:
        tree.right = insert(tree.right, val)
    return tree


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


def bst_distance(nums: List[int], a: int, b: int) -> int:
    '''
    >>> bst_distance([2, 1, 3], 1, 3)
    2
    >>> bst_distance([2, 1, 3], 1, 4)
    -1
    >>> bst_distance([2, 1, 3], 2, 1)
    1
    >>> bst_distance([2, 1, 3], 3, 3)
    0
    '''
    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))