kdtree

import%20marimo%0A%0A__generated_with%20%3D%20%220.24.0%22%0Aapp%20%3D%20marimo.App()%0A%0A%0A%40app.cell%0Adef%20_()%3A%0A%20%20%20%20import%20marimo%20as%20mo%0A%0A%20%20%20%20return%20(mo%2C)%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20%3Ca%20href%3D%22https%3A%2F%2Fcolab.research.google.com%2Fgithub%2Fapplejxd%2Fcolaboratory%2Fblob%2Fmaster%2Falgorithm%2Fkdtree.ipynb%22%20target%3D%22_parent%22%3E%3Cimg%20src%3D%22https%3A%2F%2Fcolab.research.google.com%2Fassets%2Fcolab-badge.svg%22%20alt%3D%22Open%20In%20Colab%22%2F%3E%3C%2Fa%3E%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_()%3A%0A%20%20%20%20from%20typing%20import%20List%2C%20Optional%0A%0A%20%20%20%20class%20KDNode%3A%0A%20%20%20%20%20%20%20%20def%20__init__(self%2C%20point%3A%20List%5Bfloat%5D%2C%20left%3DNone%2C%20right%3DNone)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20self.point%3A%20List%5Bfloat%5D%20%3D%20point%0A%20%20%20%20%20%20%20%20%20%20%20%20self.left%3A%20Optional%5BKDNode%5D%20%3D%20left%0A%20%20%20%20%20%20%20%20%20%20%20%20self.right%3A%20Optional%5BKDNode%5D%20%3D%20right%0A%0A%20%20%20%20return%20KDNode%2C%20List%2C%20Optional%0A%0A%0A%40app.cell%0Adef%20_(KDNode%2C%20List%2C%20Optional)%3A%0A%20%20%20%20import%20matplotlib.pyplot%20as%20plt%0A%0A%20%20%20%20class%20KDTree%3A%0A%20%20%20%20%20%20%20%20def%20__init__(self%2C%20points%3A%20List%5BList%5Bfloat%5D%5D)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20self.root%20%3D%20self.build_tree(points)%0A%0A%20%20%20%20%20%20%20%20def%20build_tree(self%2C%20points%3A%20List%5BList%5Bfloat%5D%5D)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20if%20not%20points%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20return%20None%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20k%20%3D%20len(points%5B0%5D)%20%20%23%20%E6%AC%A1%E5%85%83%E6%95%B0%0A%20%20%20%20%20%20%20%20%20%20%20%20stack%20%3D%20%5B(points%2C%200%2C%20None%2C%20None)%5D%20%20%23%20(%E7%82%B9%E3%81%AE%E3%83%AA%E3%82%B9%E3%83%88%2C%20%E6%AC%A1%E5%85%83%2C%20%E8%A6%AA%E3%81%AE%E3%83%8E%E3%83%BC%E3%83%89%2C%20%E4%BD%8D%E7%BD%AE)%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20root%3A%20Optional%5BKDNode%5D%20%3D%20None%0A%20%20%20%20%20%20%20%20%20%20%20%20while%20stack%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20pts%2C%20depth%2C%20parent%2C%20position%20%3D%20stack.pop()%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20if%20not%20pts%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20continue%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20axis%20%3D%20depth%20%25%20k%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20pts.sort(key%3Dlambda%20x%3A%20x%5Baxis%5D)%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20median%20%3D%20len(pts)%20%2F%2F%202%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20node%20%3D%20KDNode(pts%5Bmedian%5D)%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20if%20not%20root%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20root%20%3D%20node%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20else%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20if%20position%20%3D%3D%20%22left%22%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20parent.left%20%3D%20node%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20elif%20position%20%3D%3D%20%22right%22%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20parent.right%20%3D%20node%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20left_points%20%3D%20pts%5B%3Amedian%5D%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20right_points%20%3D%20pts%5Bmedian%2B1%3A%5D%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20stack.append((right_points%2C%20depth%20%2B%201%2C%20node%2C%20%22right%22))%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20stack.append((left_points%2C%20depth%20%2B%201%2C%20node%2C%20%22left%22))%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20return%20root%0A%0A%20%20%20%20%20%20%20%20def%20print_tree(self%2C%20node%2C%20depth%3D0)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20if%20not%20node%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20return%0A%20%20%20%20%20%20%20%20%20%20%20%20print(%22%20%20%22%20*%20depth%20%2B%20str(node.point))%0A%20%20%20%20%20%20%20%20%20%20%20%20self.print_tree(node.left%2C%20depth%20%2B%201)%0A%20%20%20%20%20%20%20%20%20%20%20%20self.print_tree(node.right%2C%20depth%20%2B%201)%0A%0A%20%20%20%20%20%20%20%20def%20plot_tree(self%2C%20ax%2C%20node%2C%20xmin%2C%20xmax%2C%20ymin%2C%20ymax%2C%20depth%3D0)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20if%20not%20node%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20return%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20k%20%3D%20len(node.point)%20%20%23%20%E6%AC%A1%E5%85%83%E6%95%B0%0A%20%20%20%20%20%20%20%20%20%20%20%20axis%20%3D%20depth%20%25%20k%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20if%20axis%20%3D%3D%200%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20ax.plot(%5Bnode.point%5B0%5D%2C%20node.point%5B0%5D%5D%2C%20%5Bymin%2C%20ymax%5D%2C%20'r-')%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20self.plot_tree(ax%2C%20node.left%2C%20xmin%2C%20node.point%5B0%5D%2C%20ymin%2C%20ymax%2C%20depth%20%2B%201)%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20self.plot_tree(ax%2C%20node.right%2C%20node.point%5B0%5D%2C%20xmax%2C%20ymin%2C%20ymax%2C%20depth%20%2B%201)%0A%20%20%20%20%20%20%20%20%20%20%20%20else%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20ax.plot(%5Bxmin%2C%20xmax%5D%2C%20%5Bnode.point%5B1%5D%2C%20node.point%5B1%5D%5D%2C%20'b-')%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20self.plot_tree(ax%2C%20node.left%2C%20xmin%2C%20xmax%2C%20ymin%2C%20node.point%5B1%5D%2C%20depth%20%2B%201)%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20self.plot_tree(ax%2C%20node.right%2C%20xmin%2C%20xmax%2C%20node.point%5B1%5D%2C%20ymax%2C%20depth%20%2B%201)%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20ax.plot(node.point%5B0%5D%2C%20node.point%5B1%5D%2C%20'ko')%0A%20%20%20%20%20%20%20%20%20%20%20%20ax.text(node.point%5B0%5D%20%2B%200.2%2C%20node.point%5B1%5D%20%2B%200.2%2C%20f'D%7Bdepth%7D'%2C%20color%3D'black'%2C%20fontsize%3D12%2C%20ha%3D'right')%0A%0A%20%20%20%20%20%20%20%20def%20plot_points(self%2C%20ax%2C%20points)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20for%20point%20in%20points%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20ax.plot(point%5B0%5D%2C%20point%5B1%5D%2C%20'go')%20%20%23%20'go'%20%E3%81%AF%E7%B7%91%E3%81%AE%E4%B8%B8%20(green%20circles)%20%E3%82%92%E6%84%8F%E5%91%B3%E3%81%97%E3%81%BE%E3%81%99%0A%0A%20%20%20%20return%20KDTree%2C%20plt%0A%0A%0A%40app.cell%0Adef%20_(KDTree)%3A%0A%20%20%20%20points%20%3D%20%5B(2%2C%203)%2C%20(5%2C%204)%2C%20(9%2C%206)%2C%20(4%2C%207)%2C%20(8%2C%201)%2C%20(7%2C%202)%5D%0A%20%20%20%20kd_tree%20%3D%20KDTree(points)%0A%0A%20%20%20%20kd_tree.print_tree(kd_tree.root)%0A%20%20%20%20return%20kd_tree%2C%20points%0A%0A%0A%40app.cell%0Adef%20_(kd_tree%2C%20plt%2C%20points)%3A%0A%20%20%20%20fig%2C%20ax%20%3D%20plt.subplots()%0A%20%20%20%20ax.set_xlim(0%2C%2010)%0A%20%20%20%20ax.set_ylim(0%2C%2010)%0A%20%20%20%20kd_tree.plot_tree(ax%2C%20kd_tree.root%2C%200%2C%2010%2C%200%2C%2010)%0A%20%20%20%20kd_tree.plot_points(ax%2C%20points)%20%20%23%20%E3%81%99%E3%81%B9%E3%81%A6%E3%81%AE%E7%82%B9%E3%82%92%E3%83%97%E3%83%AD%E3%83%83%E3%83%88%0A%20%20%20%20plt.show()%0A%20%20%20%20return%0A%0A%0Aif%20__name__%20%3D%3D%20%22__main__%22%3A%0A%20%20%20%20app.run()%0A