Unityで点群データの効率的な探索をするために3次元のkd木探索法が使えるかなと思ったのですが、C#で書く前にグラフを出力して確認しながらアルゴリズムを理解した方が良いと思ったので一度Pythonで実装してみます。

....簡単に実装したいなら既にscipyなどのライブラリで提供されているのでそちらを使うと良いと思います。

SciPyのKDTreeを試す

参考

以下の@fj-thさんの記事を読んでいることを前提とします。
アルゴリズムの説明とC言語でkd-treeの実装をしています。

2次元のkd木

3次元のkd木は、2次元のkd木を拡張するれば良いので、まず@fj-thさんのC言語のkd-木(リポジトリ)をPythonに書き換えます。
ただPythonで書き換えただけなので説明は割愛します。
以下のノートブックを参照ください。

kd_tree_2d.ipynb

ノードによる分割結果

3次元のkd木

kd-treeクラス

2次元のkd木を拡張する形で実装します。

kd_tree_3d.ipynb

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
"""
特定列を基準にソート
axis_xy : xなら0 yなら1 zなら2
"""
def sortxyz(arr:np.ndarray,axis_xyz:int,offset:int = 0):
    axis = axis_xyz + offset # offset
    return arr[np.argsort(arr[:,axis])]

# test
arr = np.array([[1000000,2,50],
               [2000000,0,100],
               [3000000,1,200]])
print(sortxyz(arr,0))

x_col = 0
y_col = 1
z_col = 2
class Node:
    def set_node(self,points:np.ndarray,right:int,depth:int):
        if right < 0:
            return None
        elif right == 0:
            return self.set_leaf(points[right],depth)
        
        axis = depth % 3
        sorted_points = sortxyz(points[:right + 1],axis)
        if axis == 0:
            self.left_most = sorted_points[0][x_col]
            self.right_most = sorted_points[right][x_col]
        elif axis == 1:
            self.bottom_most = sorted_points[0][y_col]
            self.top_most = sorted_points[right][y_col]
        else:
            self.front_most = sorted_points[0][z_col]
            self.back_most = sorted_points[right][z_col]
            
        median = int(right / 2)
        self.border = (sorted_points[median][x_col:] + sorted_points[median+1][x_col:])/2.0
        self.location = sorted_points[median]
        self.depth = depth
        self.right_child = Node().set_node(sorted_points[median
                                                          + 1:],right -( median + 1),depth + 1)
        self.left_child = Node().set_node(sorted_points,median,depth + 1)
        
        #ここから先は関連付け
        if not(axis == 1):
            if(not(self.right_child is None) and not(self.left_child is None)):
                self.top_most = self.right_child.top_most if self.right_child.top_most > self.left_child.top_most else self.left_child.top_most
                self.bottom_most = self.right_child.bottom_most if self.right_child.bottom_most < self.left_child.bottom_most else self.left_child.bottom_most
            elif not(self.right_child is None):
                self.top_most = self.right_child.top_most
                self.bottom_most = self.right_child.bottom_most
            elif not(self.left_child is None):
                self.top_most = self.left_child.top_most
                self.bottom_most = self.left_child.bottom_most
            else:
                self.top_most = self.location[y_col]
                self.bottom_most = self.location[y_col]
        if not(axis == 0):
            if(not(self.right_child is None) and not(self.left_child is None)):
                self.right_most = self.right_child.right_most if self.right_child.right_most > self.left_child.right_most else self.left_child.right_most
                self.left_most = self.right_child.left_most if self.right_child.left_most < self.left_child.left_most else self.left_child.left_most
            elif not(self.right_child is None):
                self.right_most = self.right_child.right_most
                self.left_most = self.right_child.left_most
            elif not(self.left_child is None):
                self.right_most = self.left_child.right_most
                self.left_most = self.left_child.left_most
            else:
                self.right_most = self.location[x_col]
                self.left_most = self.location[x_col]
        if not(axis == 2):
            if(not(self.right_child is None) and not(self.left_child is None)):
                self.back_most = self.right_child.back_most if self.right_child.back_most > self.left_child.back_most else self.left_child.back_most
                self.front_most = self.right_child.front_most if self.right_child.front_most < self.left_child.front_most else self.left_child.front_most
            elif not(self.right_child is None):
                self.back_most = self.right_child.back_most
                self.front_most = self.right_child.front_most
            elif not(self.left_child is None):
                self.back_most = self.left_child.back_most
                self.front_most = self.left_child.front_most
            else:
                self.back_most = self.location[z_col]
                self.front_most = self.location[z_col]
                
        return self
    def set_leaf(self,location:np.ndarray,depth:int):
        self.location = location
        self.left_child = None
        self.right_child = None
        self.depth = depth
        self.left_most = location[x_col]
        self.right_most = location[x_col]
        self.top_most =   location[y_col]
        self.bottom_most = location[y_col]
        self.front_most = location[z_col]
        self.back_most = location[z_col]
        return self
    def is_contained(self,sx:int,tx:int,sy:int,ty:int,sz:int,tz:int):
        """
        > regionノードから始まるsubtreeの要素が指定領域にすっぽり収まっているか否かを返す。  
        > 各ノードに、自分以下の子の最大、最小値をもたせているので、それを参照するだけで判定できる。
        実装参考のis_contained関数より引用
        """
        return not(self.left_most < sx or self.right_most > tx or \
                    self.top_most > ty or self.bottom_most < sy or \
                    self.front_most < sz or self.back_most > tz)
class Tree:
    def build(self,points):
        node = Node().set_node(points,len(points) - 1,0)
        self.top_node = node
        return node
    def search(self,r:list):
        x,y,z = r
        sx,tx = x
        sy,ty = y
        sz,tz = z
        search_results = []
        def _search(v:Node):
            nonlocal search_results,sx,tx,sy,ty
            if v.right_most < sx or v.left_most > tx or \
                v.bottom_most > ty or v.top_most < sy or \
                v.front_most > tz or v.back_most < sz:
                return;
            if v.left_child is None and v.right_child is None:
                if sx <= v.location[x_col] and sy <= v.location[y_col] and  sz <= v.location[z_col] and\
                    tx >= v.location[x_col] and ty >= v.location[y_col] and  tz >= v.location[z_col]:
                    self.search_results.append(v.location)
                    return
            if not(v.left_child is None):
                if(v.left_child.is_contained(sx,tx,sy,ty,sz,tz)):
                    search_results += self.report_subtree(v.left_child)
                else:
                    _search(v.left_child)
            if not(v.right_child is None):
                if(v.right_child.is_contained(sx,tx,sy,ty,sz,tz)):
                    search_results += self.report_subtree(v.right_child)
                else:
                    _search(v.right_child)
        _search(self.top_node)
        return search_results
    def report_subtree(self,node:Node = None,parent:Node=None,parent_border_axis:tuple = None,draw_border:tuple = None):
        """
        draw_border:tuple = (ax,max_x,min_x,max_y,min_y,max_z,min_z)
        """
        if node is None:
            node = self.top_node
        if node.left_child is None and node.right_child is None:
            return [node.location]
        pba = parent_border_axis
        if not (draw_border is None) :
            def _plot(axis:int,s,e):
                if axis == 0:
                    x = [s[0],e[0]]
                    y = [s[1],e[1]]
                    z = np.linspace(s[2], e[2])
                    Y, Z = np.meshgrid(y, z)
                    X = np.array([x] * Y.shape[0])
                    ax.plot_surface(X, Y, Z) 
                elif axis == 1:
                    x = [s[0],e[0]]
                    y = [s[1],e[1]]
                    z = np.linspace(s[2], e[2])
                    X, Z = np.meshgrid(x, z)
                    Y = np.array([y] * X.shape[0])
                    ax.plot_surface(X, Y, Z) 
                else:
                    x = np.linspace(s[0], e[0])
                    y = [s[1],e[1]]
                    z = [s[2],e[2]]
                    Z, X = np.meshgrid(z,x)
                    Y = np.array([y] * Z.shape[0])
                    ax.plot_surface(X, Y, Z)
            ax,xmax,xmin,ymax,ymin,zmax,zmin = draw_border
            s = node.border.copy()
            e = [0,0,0]
            if parent is None:
                s[1] = ymax
                s[2] = zmax
                e = np.array([s[0],ymin,zmin])
                pba = (s[0],None,s[2])
                _plot(0,s,e)
            else:
                if not(pba[0] is None) and not(pba[2] is None): #前がxの区切り(yz線) => 今がyの区切り(xz線)
                    x = xmax if(pba[0] < s[0]) else xmin
                    z = zmax if(pba[2] < s[2]) else zmin
                    s[0] = pba[0]
                    s[2] = pba[2]
                    e = np.array([x,s[1],z])
                    pba = (s[0],s[1],None)
                    _plot(1,s,e)
                elif not(pba[1] is None) and not(pba[0] is None): #前がyの区切り(xz線) => 今がzの区切り(xy線)

                    x = xmax if(pba[0] < s[0]) else xmin
                    y = ymax if(pba[1] < s[1]) else ymin
                    s[1] = pba[1]
                    s[0] = pba[0]
                    e = np.array([x,y,s[2]])
                    pba = (None,s[1],s[2])
                    _plot(2,s,e)
                else: #前がzの区切り(xy線) => 今がxの区切り(yz線)
                    y = ymax if(pba[1] < s[1]) else ymin
                    z = zmax if(pba[2] < s[2]) else zmin
                    s[2] = pba[2]
                    s[0] = pba[0]
                    e = np.array([s[0],y,z])
                    pba = (s[0],None,s[2])
                    _plot(0,s,e)
            
        arr = []
        if not (node.left_child is None):
            arr += self.report_subtree(node.left_child,node,pba,draw_border)
        if not (node.right_child is None):
            arr += self.report_subtree(node.right_child,node,pba,draw_border)
        return arr

追加したNodeクラスのfront_mostとback_mostが各ノードのz軸の範囲を持ちます。

テストデータ作成とkd-tree構築

テストデータを使って境界を表示してみます。

# test
points = [[1,2,2],
          [2,4,4],
          [5,6,3],
          [7,8,6],
          [3,10,4],
          [11,11,3],
         [10,10,4],
         [0,8,7]]
points = np.array(points)
tree = Tree()
node = tree.build(points)
fig = plt.figure()
ax = fig.add_subplot(projection='3d')
ax.scatter(points[:,x_col],points[:,y_col],points[:,z_col])
plt.show()
%matplotlib notebook
fig = plt.figure()
ax = fig.gca(projection='3d')
ax.scatter(points[:,x_col],points[:,y_col],points[:,z_col])
tree.report_subtree(draw_border=(ax,points[:,x_col].max(),points[:,x_col].min(),points[:,y_col].max(),points[:,y_col].min(),points[:,z_col].max(),points[:,z_col].min()))
#ax.view_init(elev=10, azim=50)
plt.show()

分割結果

グルグルグラフを回転させてみると点が分けられていることがわかると思います。

探索

次に探索してみます。search関数の引数にx,y,zの範囲を指定すると返り値にその範囲内の点が返されます。

results = np.array(tree.search([[0,10],[8,10],[0,6]]))
print(results)
fig = plt.figure()
ax = fig.gca(projection='3d')
ax.scatter(results[:,x_col],results[:,y_col],results[:,z_col])
plt.xlim([points[:,x_col].min()-1,points[:,x_col].max()+1])
plt.ylim([points[:,y_col].min()-1,points[:,y_col].max()+1])
plt.show()

"""
results : [[ 3 10  4]
 [ 7  8  6]
 [10 10  4]]
"""

(グラフがわかりにくいですが)できました。
どちらかというと3次元に拡張するより、アルゴリズムの理解と境界のグラフ表示に時間がかかりました。