关于KDTree

        K-D Tree(KDT, K-Dimension Tree)是一种可以高效处理K维空间信息的数据结构。在结点数n远大于\(2^k\)时，应用K-D Tree的时间效率很好(载自OI Wiki)。

环境

        Windows10         Unity2022.3.34f1c1

建树

        在一堆数据中，我们先选取一个维度，在这个维度上我们遍历数据选择出一个中间值。以这个中间值为界将剩下的数据分成两个部分：比其大在右边，小于等于其的在左边。在遍历的过程中，我们可以得到一个这堆数据中在各个维度上的最大值和最小值。我们也要将这信息记录在节点中。接下来我们选择另一个维度，重复上述过程，如果维度用完了就再从头按照之前的顺序重复过程，直到最终叶子节点记录所有的数据。非叶子节点本身是仅记录其对应数据的范围，即我们在遍历过程中得到的各个维度上的最大值和最小值。关于维度，你可以想你要对一个三维空间下的物体进行划分，它有一个位置信息：x,y,z。那么我们就可以从x->y->z的顺序选择维度。当然顺序的选择看你如果操作。实际操作过程中，你也不一定要做维度的轮换，反正只要按照一定的规则去选择维度进行划分就好。我认为做维度分割是为了建立出来的树更平衡，方便后面的搜索，所以只要选择适合的方式去做维度选择就好。

        在实际操作的过程中，不一定要叶子节点才存储数据。你也可以在用树的结点去做存储。叶子节点也不一定只存储一个数据，它也可以存储多个。维度的选择也不是固定的，实际上为了更好的搜索，应该判断一下那个维度分离的数据数量差值最少。比如你发现使用x划分数据得到的两个数据堆数量为16,15，而使用y划分则是1,30。那么即使你之前使用了x做划分了，为了高效搜索，你这样应该使用y做划分。很多资料中都显示说要使用中位数，我想本身这也是为了分割的数据集更均匀。

搜索

        如果我们要查找一个范围内的数据。首先我们用根节点判断。在建树的时候，我们存储下来每个节点的数据范围，我们判断一下这个搜索范围是否和我们的数据范围有交集。如果没有交集，那就说明这个范围内没有数据，我们直接返回空。如果有交集，我们则对其子节点重复上述的操作。直到到叶子节点为止。

实际操作

        KDTree的实现有挺多的，我这边就做一个k=2的例子，其中用叶子节点存储和不用其存储我都会做一遍。这里我推荐一个Unity官方提供的kdtree实现https://github.com/Unity-Technologies/Megacity-2019/blob/master/Assets/Scripts/Utils/KDTree/KDTree/KDTree.cs。它使用了并行计算，并且支持多线程。它的实现就是用叶子节点存储数据。我更希望你可以看看它的实现，因为我自己并不能我的实现一定对。但是Unity官方这个实现，我自己项目也在用，而且也经过了许多测试。我个人认为它的实现很好的。

        叶子节点存储数据，我想仿造Unity官方提供的kdtree实现，使用数组来存储数据而不是用一个树结构存储。而树节点存储数据，我就使用常用的树结构来存储数据了。

叶子节点存储数据

        因为我们使用数组来存储树信息，且我们使用叶子节点来保存真正的数据。又因为我们要知道每一个节点所包围的范围，因此叶子节点也需要存储其所函数据的范围。故节点数据如下设定：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
// 设定的范围是方形包围盒，所以只需要最大值和最小值。
public struct Bounds
{
    // 最小点
    public Vector3 min;
    // 最大点
    public Vector3 max;
}

public class TreeNode
{
    // 是否是叶子节点
    public bool isLeaf;

    // 包围盒
    public Bounds bounds;

    // 数据下标
    public int dataInd = -1;
}

这里我虽然使用 Vector3 来做存储，但实际上我还是只用了 x 和 y 两个维度。接下来便是kdtree的数据设定。因为这里我是仿造Unity官方提供的kdtree去做的实现，所以这里我需要用额外的信息记录下传进来的数据。我将这个额外的信息命名为KDTreeData。

1
2
3
4
5
6
7
public struct KDTreeData
{
    // 原始数据中对应的下标
    public int ind;
    // 所对应的数据
    public Vector3 val;
}

kdtree的数据结构如下：

1
2
3
4
5
6
public class KDTreeByLeaf
{
    private List<KDTreeData> data;

    private List<TreeNode> nodes;
}

其中nodes是树的节点信息，而data是原始数据的备份。结构确定后，我们便可以开始实现kdtree的建树过程。因为我们使用数组来做树的存储，所以我们要一开始就知道我们需要开大多的数组空间去存储树结构。我们使用叶子节点来存储数据，且我们又要尽量保证实现的树是一个平衡树。既然如此，那么我们只要得到大于数据数量的最小2幂的值。然后我们通过最大叶子节点的数量来倒推出总的节点数量。公式即叶子节点的数量 * 2 - 1 = 总节点的数量。接下来就是按照建树的想法去做操作。这里我偷懒使用了快排来得到中间值，如果为了追求效率大家可以换成查找中间值的算法比如BFPTR算法。不过我因为偷懒使用了快排，所以我又多实现了一些快排的比较操作。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
public class KDTreeXCompare : IComparer<KDTreeData>
{
    public int Compare(KDTreeData x, KDTreeData y)
    {
        return x.val.x.CompareTo(y.val.x);
    }
}

public class KDTreeYCompare : IComparer<KDTreeData>
{
    public int Compare(KDTreeData x, KDTreeData y)
    {
        return x.val.y.CompareTo(y.val.y);
    }
}

下面是完整的建树过程：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
public void BuildTree(List<Vector3> positions)
{
    // 先复制一份数据
    data = new List<KDTreeData>();
    for (int i = 0;i < positions.Count; ++i)
    {
        data.Add(new KDTreeData()
        {
            ind = i,
            val = positions[i]
        });
    }

    // 创建节点，因为我们用叶子节点来存储数据，所以我们节点数量会多于数据数量
    nodes = new List<TreeNode>();
    int nodeNum = 1;
    while(nodeNum < positions.Count)
    {
        nodeNum *= 2;
    }
    nodeNum = nodeNum * 2 - 1;
    for (int i = 0; i < nodeNum; ++i)
    {
        nodes.Add(new TreeNode());
    }
    BuildSubTree(0, positions.Count, 0);
}

 private void BuildSubTree(int startInd,int endInd,int nodeInd)
 {
     var curNode = nodes[nodeInd];
     if(endInd - startInd == 1)
     {
         curNode.isLeaf = true;
         curNode.dataInd = startInd;
         return;
     }

     // 获取范围
     Vector3 min = new Vector3(float.MaxValue,float.MaxValue, float.MaxValue);
     Vector3 max = new Vector3(float.MinValue, float.MinValue, float.MinValue);

     for (int i = startInd;i < endInd; ++i)
     {
         var info = data[i].val;
         min = Vector3.Min(min, info);
         max = Vector3.Max(max, info);
     }

     curNode.bounds = new Bounds
     {
         min = min,
         max = max,
     };

     // 决定用那个方向进行分割
     float lengthX = max.x - min.x;
     float lengthY = max.y - min.y;

     if(lengthX > lengthY)
         data.Sort(startInd, endInd - startInd, new KDTreeXCompare());
     else
         data.Sort(startInd, endInd - startInd, new KDTreeYCompare());

     // 因为排序了，所以这里直接二分就好了
     int leftInd = startInd + (endInd - startInd) / 2;

     BuildSubTree(startInd, leftInd, nodeInd * 2 + 1);
     BuildSubTree(leftInd, endInd, nodeInd * 2 + 2);
 }

在实现搜索的时候，我们需要注意我们返回的应该是原数组的下标。这里判断两个矩形是否相交，我借鉴了这篇【数据结构和算法】判断两个矩形是否相交文章的思想。具体实现如下:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
public List<int> Search(Vector3 pos,float radius)
{
    List<int> result = new List<int>();
    Queue<int> q = new Queue<int>();
    q.Enqueue(0);
    while (q.Count > 0)
    {
        var ind = q.Dequeue();
        var curNode = nodes[ind];
        if(curNode.isLeaf)
        {
            var curInfo = data[curNode.dataInd];
            if(Vector2.Distance(curInfo.val,pos) <= radius)
            {
                result.Add(curInfo.ind);
            }
        }
        else
        {
            // 判断是否在包围盒内
            var center = (curNode.bounds.min + curNode.bounds.max) / 2;
            var xval = center.x - pos.x;
            xval = center.x - curNode.bounds.min.x + radius - xval;
            var yval = center.y - pos.y;
            yval = center.y - curNode.bounds.min.y + radius - yval;
            if(xval > 0 && yval > 0)
            {
                q.Enqueue(ind * 2 + 1);
                q.Enqueue(ind * 2 + 2);
            }
        }
    }
    return result;
}

最后关于KDTreeByLeaf的完整实现如下，下面代码中没有在上文出现的函数实现或变量都是我用来做验证的，有些涉及到了可视化，其中还多包含了一些打印信息。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
public class KDTreeByLeaf : IDisposable
{
    private List<KDTreeData> data;

    private List<TreeNode> nodes;

    private List<Vector2> bounds;

    public void BuildTree(List<Vector3> positions)
    {
        // 先复制一份数据
        data = new List<KDTreeData>();
        for (int i = 0;i < positions.Count; ++i)
        {
            data.Add(new KDTreeData()
            {
                ind = i,
                val = positions[i]
            });
        }

        // 创建节点，因为我们用叶子节点来存储数据，所以我们节点数量会多于数据数量
        nodes = new List<TreeNode>();
        int nodeNum = 1;
        while(nodeNum < positions.Count)
        {
            nodeNum *= 2;
        }
        nodeNum = nodeNum * 2 - 1;
        for (int i = 0; i < nodeNum; ++i)
        {
            nodes.Add(new TreeNode());
        }
        BuildSubTree(0, positions.Count, 0);
    }
    public void Dispose()
    {
        data.Clear();
    }

    public void DebugInfo()
    {
        foreach (TreeNode node in nodes)
        {
            if(node.isLeaf)
            {
                Debug.Log(node.dataInd);
            }
        }
    }

    public void ShowTree(Vector3 centerPos)
    {
        if (nodes == null || nodes.Count == 0)
            return;
        //Debug.Log(nodes.Count);
        CalcTreeNodeShowBound();
        ShowNodeInfo(centerPos, 0);
    }

    public void ShowNodeInfo(Vector3 centerPos,int ind)
    {
        if (nodes == null || nodes.Count < ind)
            return;

        var nodeInfo = nodes[ind];
        if(nodeInfo.isLeaf)
        {
            if(nodeInfo.dataInd != -1)
                Gizmos.color = Color.red;
            else
                Gizmos.color = Color.grey;
        }
        else
        {
            var curBounds = bounds[ind] * 0.25f;
            Vector3 leftPos = centerPos + new Vector3(-1 * curBounds.x, -1, 0);
            Vector3 right = centerPos + new Vector3(curBounds.y , -1, 0);
            Gizmos.color = Color.black;
            Gizmos.DrawLine(leftPos, centerPos);
            Gizmos.DrawLine(right, centerPos);
            ShowNodeInfo(leftPos, ind * 2 + 1);
            ShowNodeInfo(right, ind * 2 + 2);
            Gizmos.color = Color.green;
        }
        Gizmos.DrawSphere(centerPos, 0.25f);
    }

    public List<int> Search(Vector3 pos,float radius)
    {
        List<int> result = new List<int>();
        Queue<int> q = new Queue<int>();
        q.Enqueue(0);
        while (q.Count > 0)
        {
            var ind = q.Dequeue();
            var curNode = nodes[ind];
            if(curNode.isLeaf)
            {
                var curInfo = data[curNode.dataInd];
                if(Vector2.Distance(curInfo.val,pos) <= radius)
                {
                    result.Add(curInfo.ind);
                }
            }
            else
            {
                // 判断是否在包围盒内
                var center = (curNode.bounds.min + curNode.bounds.max) / 2;
                var xval = center.x - pos.x;
                xval = center.x - curNode.bounds.min.x + radius - xval;
                var yval = center.y - pos.y;
                yval = center.y - curNode.bounds.min.y + radius - yval;
                if(xval > 0 && yval > 0)
                {
                    q.Enqueue(ind * 2 + 1);
                    q.Enqueue(ind * 2 + 2);
                }
            }
        }
        return result;
    }

    private void CalcTreeNodeShowBound()
    {
        if(bounds == null || bounds.Count != nodes.Count)
        {
            bounds = new List<Vector2>();
            foreach (var i in nodes)
            {
                bounds.Add(Vector2.zero);
            }
            PreCalc(0);
        }
    }

    private void PreCalc(int ind)
    {
        var nodeInfo = nodes[ind];
        if(nodeInfo.isLeaf)
        {
            if (ind % 2 == 1)
                bounds[ind] = Vector2.right;
            else bounds[ind] = Vector2.up;
            return;
        }
        PreCalc(ind * 2 + 1);
        var leftVal = bounds[ind * 2 + 1];
        PreCalc(ind * 2 + 2);
        var rightVal = bounds[ind * 2 + 2];
        var curVal = new Vector2(leftVal.x + leftVal.y + 1, rightVal.x + rightVal.y + 1);
        bounds[ind] = curVal;
    }

    private void BuildSubTree(int startInd,int endInd,int nodeInd)
    {
        var curNode = nodes[nodeInd];
        if(endInd - startInd == 1)
        {
            if (curNode.isLeaf)
                Debug.Log("Repeat");

            curNode.isLeaf = true;
            curNode.dataInd = startInd;
            return;
        }

        // 获取范围，平均值，中间中
        Vector3 min = new Vector3(float.MaxValue,float.MaxValue, float.MaxValue);
        Vector3 max = new Vector3(float.MinValue, float.MinValue, float.MinValue);

        for (int i = startInd;i < endInd; ++i)
        {
            var info = data[i].val;
            min = Vector3.Min(min, info);
            max = Vector3.Max(max, info);
        }

        curNode.bounds = new Bounds
        {
            min = min,
            max = max,
        };

        // 决定用那个方向进行分割
        float lengthX = max.x - min.x;
        float lengthY = max.y - min.y;

        if(lengthX > lengthY)
            data.Sort(startInd, endInd - startInd, new KDTreeXCompare());
        else
            data.Sort(startInd, endInd - startInd, new KDTreeYCompare());
        int leftInd = startInd + (endInd - startInd) / 2;

        BuildSubTree(startInd, leftInd, nodeInd * 2 + 1);
        BuildSubTree(leftInd, endInd, nodeInd * 2 + 2);
    }
}

树节点存储数据

        这部分展示如何使用树节点存储数据，这部分我用传统的树结构来做对应的实现，并实现了之前的说明的分割维度轮换机制。首先是节点数据如下设定：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
public class TreeNodeData
{
    public bool isLeaf
    {
        get
        {
            return left == null && right == null;
        }
    }

    public GameObject data;

    public Bounds bounds;

    // 用来做展示用的
    public Vector2 showBounds;

    public TreeNodeData left;
    public TreeNodeData right;
}

代码中的showBounds是用来做其他功能的，可以忽略。其实剩下的也没什么特别，所以我就直接展示kdtree完整实现如下：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
public class KDTreeObjXCompare : IComparer<GameObject>
{
    public int Compare(GameObject x, GameObject y)
    {
        return x.transform.position.x.CompareTo(y.transform.position.x);
    }
}

public class KDTreeObjYCompare : IComparer<GameObject>
{
    public int Compare(GameObject x, GameObject y)
    {
        return x.transform.position.y.CompareTo(y.transform.position.y);
    }
}

public class KDTreeByNode
{
    public TreeNodeData root;

    private List<Vector2> bounds;

    public void BuildTree(List<GameObject> data)
    {
        root  = new TreeNodeData();
        BuildSubTree(root, data, 0, data.Count, 0);
    }

    private void BuildSubTree(TreeNodeData curNode, List<GameObject> data, int startInd, int endInd ,int flag)
    {
        if (endInd - startInd == 1)
        {
            curNode.data = data[startInd];
            return;
        }

        // 获取范围
        Vector3 min = new Vector3(float.MaxValue, float.MaxValue, float.MaxValue);
        Vector3 max = new Vector3(float.MinValue, float.MinValue, float.MinValue);

        for (int i = startInd; i < endInd; ++i)
        {
            var info = data[i].transform.position;
            min = Vector3.Min(min, info);
            max = Vector3.Max(max, info);
        }

        curNode.bounds = new Bounds
        {
            min = min,
            max = max,
        };

        int nextFlag = (flag + 1) % 2;

        if(flag == 0)
            data.Sort(startInd, endInd - startInd, new KDTreeObjXCompare());
        else
            data.Sort(startInd, endInd - startInd, new KDTreeObjYCompare());

        int leftInd = startInd + (endInd - startInd) / 2;
        curNode.data = data[leftInd];

        if (leftInd - startInd > 0)
        {
            curNode.left = new TreeNodeData();
            BuildSubTree(curNode.left, data, startInd, leftInd, nextFlag);
        }

        leftInd++;
        if (endInd - leftInd > 0)
        {
            curNode.right = new TreeNodeData();
            BuildSubTree(curNode.right, data, leftInd, endInd, nextFlag);
        }
    }

    public List<GameObject> Search(Vector3 pos, float radius)
    {
        List<GameObject> result = new List<GameObject>();
        Queue<TreeNodeData> q = new Queue<TreeNodeData>();
        q.Enqueue(root);
        while (q.Count > 0)
        {
            TreeNodeData curNode = q.Dequeue();
            if(curNode != null)
            {
                if (Vector2.Distance(curNode.data.transform.position, pos) <= radius)
                {
                    result.Add(curNode.data);
                }

                // 叶子节点就不判断范围了
                if (!curNode.isLeaf)
                {
                    var center = (curNode.bounds.min + curNode.bounds.max) / 2;
                    var xval = center.x - pos.x;
                    xval = center.x - curNode.bounds.min.x + radius - xval;
                    var yval = center.y - pos.y;
                    yval = center.y - curNode.bounds.min.y + radius - yval;
                    if (xval > 0 && yval > 0)
                    {
                        q.Enqueue(curNode.left);
                        q.Enqueue(curNode.right);
                    }
                }
            }
        }
        return result;
    }

    public void ShowTree(Vector3 centerPos)
    {
        if (root == null)
            return;

        CalcTreeNodeShowBound();
        ShowNodeInfo(centerPos, root);
    }

    public void ShowNodeInfo(Vector3 centerPos, TreeNodeData nodeInfo)
    {
        if(nodeInfo == null)
            Gizmos.color = Color.grey;
        else if (nodeInfo.isLeaf)
        {
            Gizmos.color = Color.red;
        }
        else
        {
            var curBounds = nodeInfo.showBounds * 0.25f;
            Vector3 leftPos = centerPos + new Vector3(-1 * curBounds.x, -1, 0);
            Vector3 rightPos = centerPos + new Vector3(curBounds.y, -1, 0);
            Gizmos.color = Color.black;
            Gizmos.DrawLine(leftPos, centerPos);
            Gizmos.DrawLine(rightPos, centerPos);
            ShowNodeInfo(leftPos, nodeInfo.left);
            ShowNodeInfo(rightPos, nodeInfo.right);
            Gizmos.color = Color.green;
        }
        Gizmos.DrawSphere(centerPos, 0.25f);
    }

    private void CalcTreeNodeShowBound()
    {
        bounds = new List<Vector2>();
        PreCalc(root);
    }

    private void PreCalc(TreeNodeData nodeInfo)
    {
        var leftVal = Vector2.right;
        if (nodeInfo.left != null)
        {
            if (nodeInfo.left.isLeaf)
                nodeInfo.left.showBounds = Vector2.right;
            else
                PreCalc(nodeInfo.left);
            leftVal = nodeInfo.left.showBounds;
        }
        var rightVal = Vector2.up;
        if (nodeInfo.right != null)
        {
            if (nodeInfo.right.isLeaf)
                nodeInfo.right.showBounds = Vector2.up;
            else
                PreCalc(nodeInfo.right);
            rightVal = nodeInfo.right.showBounds;
        }
        var curVal = new Vector2(leftVal.x + leftVal.y + 1, rightVal.x + rightVal.y + 1);
        nodeInfo.showBounds = curVal;
    }
}

参考文章

闲言碎语

        这篇文章我才写一点就知道这文章有点水了。果然算法或是数据结构的文章就真的可以很水的，而且有时候还无法验证自己是否错了。不过水就水了，至少我现在算是掌握了一点关于kdtree的知识。

        不过让我没想到的是最麻烦的竟然是后面实现的部分。一开始我想着仿造Unity官方的实现，但是我发现官方为了追求效率其中的实现极度复杂。我阅读了半天，我也不明白它是如何尽量保证自己的树是平衡树的。我删减的了部分并进行实践后发现，按照官方的划分方法的确会产生大多的节点生成在一边。特别是大部分数据紧密堆积，小部分数据极度分散的时候尤为明显。不过考虑到其实它一个叶子节点中可以存储不只一个数据，或许这就是为什么它可以正常运行的原因吧。不过我确实还是没读懂，或许它其实还是有操作来保证平衡呢？