kdtreeanim.py

from rp import *
class KDTreeNode:
    def __init__(self,points,horz_split=True,level=1,parent=None):
        
        self.points    =as_complex_vector(points)
        self.horz_split=horz_split
        self.level     =level
        
        self.median=median(map(self.key,self.points))
        self.upper_points =[point for point in points if self.key(point) >self.median]
        self.lower_points =[point for point in points if self.key(point) <self.median]
        self.middle_points=[point for point in points if self.key(point)==self.median]
        
        self.upper_node =KDTreeNode(self.upper_points ,not self.horz_split,self.level+1,self) if len(self.upper_points )   else None
        self.lower_node =KDTreeNode(self.lower_points ,not self.horz_split,self.level+1,self) if len(self.lower_points )   else None
        self.middle_node=KDTreeNode(self.middle_points,not self.horz_split,self.level+1,self) if len(self.middle_points)>1 else None
        
        self.leaf=self.middle_points[0] if (self.upper_node is None and self.lower_node is None and self.middle_node is None) else None
        
    def key(self,point):
        return point.imag if self.horz_split else point.real
    
    def points_in_rectangle(bounds):
        [x0,y0],[x1,y1]=as_points_array(bounds)
        x0,x1=sorted([x0,x1])
        y0,y1=sorted([y0,y1])
        
    
    def plot(self):
        if self.leaf:
            display_dot(self.leaf,color='red')
        else:
            if self.horz_split:
                point_a=[max(point.real for point in self.points),self.median]
                point_b=[min(point.real for point in self.points),self.median]
            else:
                point_a=[self.median,max(point.imag for point in self.points)]
                point_b=[self.median,min(point.imag for point in self.points)]
            
            display_path([point_a,point_b],color='blue',linewidth=3/self.level,alpha=1/(self.level/2))
                
            if self.upper_node :self.upper_node .plot()
            if self.lower_node :self.lower_node .plot()
            if self.middle_node:self.middle_node.plot()
    
display_clear()
KDTreeNode(random_floats_complex(100)).plot()