Pizza Delevery kd-tree k nearest neighbour search (K log k log N) avg worst k _/N got tle for some part of subtasks

for every query i find k+1 nearest neighbour and using kd-tree than ans query. but still got tle.
suggestion for any improvement in solution are welcome.

K-d trees can become unbalanced hence their worst case complexity for querying and insertion can become linear

but since i split node at median using nth element so k-d tree remain balanced and worst complexity will be _/N.