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))
Loading full content...