import numpy as np
from scipy.spatial import KDTree