如何在UE4中实现一个高效的Kd-Tree 搜索算法

Kd-Tree 有比较多的应用，在指定位置处，搜索附件圆形区域内的点集。

可以用于高效率实现非常多的跟游戏开发逻辑相关的内容，这里实现了一个高效率的基于UE4引擎的 C++ Kd-Tree搜索算法，适合性能优化的同学使用，具体使用方法 代码中有注释。

具体逻辑可以参照以下代码：

//.h 头文件 可以命名为 KdTreeContextUtils.h--------------
#include "CoreMinimal.h"
#include "HAL/ThreadSingleton.h"
#include "Components/ActorComponent.h"

//变换节点的结构体,此结构体可以进行任意扩展成需要的数据结构
struct FTransformNode
{
public:
	FVector Position;   //用于存放位置信息
	int32   Id;			//用于存放对象的实例Id

	FTransformNode() {};  

	//注意 UE 默认支持Move 拷贝构造函数即编译器默认会加入 FTransformNode&& 的构造函数  可以直接使用Move进行构造 
	FTransformNode(int32 Id) :Id(Id) {};
	FTransformNode(int32 Id, FVector Pos)
		: Id(Id)
		, Position(Pos)
	{}

};

//UE4 的线程单例
class  FContextUtils :public TThreadSingleton<FContextUtils>
{
	friend TThreadSingleton<FContextUtils>;

	FContextUtils();

	~FContextUtils();

	class SearchNode
	{
	public:
		int32 Start;
		int32 End;
		int32 LNode;
		int32 RNode;
		FVector MaxPos;
		FVector MinPos;
	};

	//递归构造搜索树
	void BuildSearchTreeRecursive(int32 Start, int32 End, int32 Node);
	//插入搜索节点
	void InsertSearchNode(const FVector& Center, float& RangeSquared, FTransformNode* SearchNode);
	//递归搜索 指定当前点和需要搜索的范围(距离的平方) 内的节点
	void QuerySearchTreeRecursive(const FVector& Center, float& RangeSquared, int32 Node);

	void BuildSearchTree();
	//需要查找的现实世界中 带有位置信息的节点
	TArray<FTransformNode> Items;
	//用于构建搜索结构的节点信息
	TArray<SearchNode> SearchTreeNodes;
	bool bRebuildTree;
}







//.cpp 文件--------------
//需要包含 #include "KdTreeContextUtils.h"
static int32 MAXLEAFSIZE = 2;
const  int32 MAXNEIGBORSIZE = 20;
struct FSearchNeighbor
{
	FTransformNode* Node = nullptr;
	float DistanceSquare;
};

struct FSearchTreeContext
{
	//存储搜索到的节点结果，可以按需搜索指定 位置 附近的距离平方的 圆形区域内的节点数量
	FSearchNeighbor Neighbors[MAXNEIGBORSIZE];
	int32 NeighborSize;


	void AddNeighbor(FTransformNode* Node, float DistanceSquare, float& RangeSquared)
	{
		if (NeighborSize < MAXNEIGBORSIZE)
		{
			Neighbors[NeighborSize].Node = Node;
			Neighbors[NeighborSize].DistanceSquare = DistanceSquare;
			++NeighborSize;
		}

		int32 i = NeighborSize - 1;
		while (0 != i && DistanceSquare < Neighbors[i - 1].DistanceSquare)
		{
			Neighbors[i] = Neighbors[i - 1];
			--i;
		}

		Neighbors[i].Node = Node;
		Neighbors[i].DistanceSquare = DistanceSquare;

		if (NeighborSize == MAXNEIGBORSIZE)
		{
			RangeSquared = Neighbors[NeighborSize - 1].DistanceSquare;
		}
	}

	void Clear()
	{
		FMemory::Memset(Neighbors, 0, MAXNEIGBORSIZE * sizeof(FSearchNeighbor));
		NeighborSize = 0;
	}
};

FSearchTreeContext SearchContext;

void FContextUtils::FContextUtils()
{
	bRebuildTree = false;
}


//构建Kd-tree
//具体步骤为  先使用数组Items 添加带有位置信息的节点
//然后递归构建Kd-tree
//QuerySearchTreeRecursive 递归查询指定的位置和范围
// 应用举例
//CurrentPos  为当前需要搜索的点的三维坐标，RangeSquared 为需要搜索的半径的平方 返回结果为查找到的区域内的所有点集的实例Id
//TArray<int32> FContextUtils::SearchNearPickUpList(FVector CurrentPos, float RangeSquared)
//{
//	TArray<int32> InstanceList;
//	BuildSearchTree();
//	SearchContext.Clear();
//	if (Items.Num())
//	{
//		QuerySearchTreeRecursive(CurrentPos, RangeSquared, 0);
//	}
//	int32 iMax = SearchContext.NeighborSize;
//	if (iMax)
//	{
//		for (int32 i = 0; i < iMax; i++)
//		{
//			int32 Id = SearchContext.Neighbors[i].Node->Id;
//			FVector PickUpPos = SearchContext.Neighbors[i].Node->Position;
//			InstanceList.Add(Id);
//		}
//	}
// return InstanceList;
//}


void FContextUtils::BuildSearchTree()
{
	if (!bRebuildTree)
	{
		return;
	}

	int32 Num = Items.Num();
	if (Num > 0)
	{	
		//采用数组方式主要是为了加速访存效率
		SearchTreeNodes.SetNum(2 * Num, false);
		BuildSearchTreeRecursive(0, Num, 0);
	}
	bRebuildTree = false;
}

void FContextUtils::BuildSearchTreeRecursive(int32 Start, int32 End, int32 Node)
{
	SearchTreeNodes[Node].Start = Start;
	SearchTreeNodes[Node].End = End;
	FVector Pos = Items[Start].Position;
	SearchTreeNodes[Node].MaxPos = Pos;
	SearchTreeNodes[Node].MinPos = Pos;
	SearchTreeNodes[Node].LNode = -1;
	SearchTreeNodes[Node].RNode = -1;

	for (int32 i = Start + 1; i < End; i++)
	{
		FVector NodePos = Items[i].Position;

		FVector MaxPos = SearchTreeNodes[Node].MaxPos;
		FVector MinPos = SearchTreeNodes[Node].MinPos;

		MaxPos[0] = FMath::Max(MaxPos[0], NodePos.X);
		MinPos[0] = FMath::Min(MinPos[0], NodePos.X);

		MaxPos[1] = FMath::Max(MaxPos[1], NodePos.Y);
		MinPos[1] = FMath::Min(MinPos[1], NodePos.Y);

		MaxPos[2] = FMath::Max(MaxPos[2], NodePos.Z);
		MinPos[2] = FMath::Min(MinPos[2], NodePos.Z);

		SearchTreeNodes[Node].MaxPos = MaxPos;
		SearchTreeNodes[Node].MinPos = MinPos;
	}

	if (End - Start > MAXLEAFSIZE)
	{
		int32 Index = 0;
		FVector MaxPos = SearchTreeNodes[Node].MaxPos;
		FVector MinPos = SearchTreeNodes[Node].MinPos;
		float Xpos = MaxPos[0] - MinPos[0];
		float Ypos = MaxPos[1] - MinPos[1];
		float Zpos = MaxPos[2] - MinPos[2];
		if (Ypos > Xpos && Ypos > Zpos)
		{
			Index = 1;
		}
		else if (Zpos > Xpos && Zpos > Ypos)
		{
			Index = 2;
		}

		const float Middle = 0.5f * (MaxPos[Index] + MinPos[Index]);

		int32 Left = Start;
		int32 Right = End;

		while (Left < Right)
		{
			while (Left < Right && Items[Left].Position[Index] < Middle)
			{
				++Left;
			}

			while (Right > Left && Items[Right - 1].Position[Index] >= Middle)
			{
				--Right;
			}

			if (Left < Right)
			{
				Swap(Items[Left], Items[Right - 1]);
				++Left;
				--Right;
			}
		}

		int32 LeftSize = Left - Start;
		if (LeftSize == 0)
		{
			++LeftSize;
			++Left;
			++Right;
		}

		int32 LNode = Node + 1;
		int32 RNode = Node + 2 * LeftSize;
		SearchTreeNodes[Node].LNode = LNode;
		SearchTreeNodes[Node].RNode = RNode;

		BuildSearchTreeRecursive(Start, Left, LNode);
		BuildSearchTreeRecursive(Left, End, RNode);
	}
}

void FContextUtils::InsertSearchNode(const FVector& Center, float& RangeSquared, FTransformNode* SearchNode)
{
	const float DistanceSquare = FVector::DistSquared(Center, SearchNode->Position);
	if (DistanceSquare < RangeSquared)
	{
		SearchContext.AddNeighbor(SearchNode, DistanceSquare, RangeSquared);
	}
}

void FContextUtils::QuerySearchTreeRecursive(const FVector& Center, float& RangeSquared, int32 Node)
{
	int32 Start = SearchTreeNodes[Node].Start;
	int32 End = SearchTreeNodes[Node].End;
	int32 Left = SearchTreeNodes[Node].LNode;
	int32 Right = SearchTreeNodes[Node].RNode;
	if (End - Start < MAXLEAFSIZE || Left < 0 || Right < 0)
	{
		for (int i = Start; i < End; i++)
		{
			InsertSearchNode(Center, RangeSquared, &Items[i]);
		}
	}
	else
	{
		const float DistSqLeft = FMath::Square(FMath::Max(0.0f, SearchTreeNodes[Left].MinPos[0] - Center.X)) +
			FMath::Square(FMath::Max(0.0f, Center.X - SearchTreeNodes[Left].MaxPos[0])) +
			FMath::Square(FMath::Max(0.0f, SearchTreeNodes[Left].MinPos[1] - Center.Y)) +
			FMath::Square(FMath::Max(0.0f, Center.Y - SearchTreeNodes[Left].MaxPos[1])) +
			FMath::Square(FMath::Max(0.0f, SearchTreeNodes[Left].MinPos[2] - Center.Z)) +
			FMath::Square(FMath::Max(0.0f, Center.Z - SearchTreeNodes[Left].MaxPos[2]));

		const float DistSqRight = FMath::Square(FMath::Max(0.0f, SearchTreeNodes[Right].MinPos[0] - Center.X)) +
			FMath::Square(FMath::Max(0.0f, Center.X - SearchTreeNodes[Right].MaxPos[0])) +
			FMath::Square(FMath::Max(0.0f, SearchTreeNodes[Right].MinPos[1] - Center.Y)) +
			FMath::Square(FMath::Max(0.0f, Center.Y - SearchTreeNodes[Right].MaxPos[1])) +
			FMath::Square(FMath::Max(0.0f, SearchTreeNodes[Right].MinPos[2] - Center.Z)) +
			FMath::Square(FMath::Max(0.0f, Center.Z - SearchTreeNodes[Right].MaxPos[2]));

		if (DistSqLeft < DistSqRight)
		{
			if (DistSqLeft < RangeSquared)
			{
				QuerySearchTreeRecursive(Center, RangeSquared, Left);

				if (DistSqRight < RangeSquared)
				{
					QuerySearchTreeRecursive(Center, RangeSquared, Right);
				}
			}
		}
		else
		{
			if (DistSqRight < RangeSquared)
			{
				QuerySearchTreeRecursive(Center, RangeSquared, Right);

				if (DistSqLeft < RangeSquared)
				{
					QuerySearchTreeRecursive(Center, RangeSquared, Left);
				}
			}
		}
	}
}