四叉树加速碰撞检测

using System.Collections.Generic;
using UnityEngine;

public class MyQuadTree
{
    public interface IShape
    {
        Rect GetBounds();
        string GetShapeName();
    }

    private Rect m_Bounds = new Rect(0, 0, 1, 1);
    private int m_Depth; //节点深度(从0开始)
    private int m_MaxDepth = 4; //树最大深度
    private MyQuadTree[] m_Children;// = new MyQuadTree[4]; //当前节点的4个子节点

    private List<IShape> m_ShapeList = new List<IShape>();
    private int m_SplitThreshold = 10; //存放的对象超过这个值时分裂为4个子节点

    public MyQuadTree(Rect bounds, int splitThreshold = 10, int maxDepth = 4)
    {
        m_Bounds = bounds;
        m_SplitThreshold = splitThreshold;
        m_MaxDepth = maxDepth;
    }

    public Rect Bounds
    {
        get { return m_Bounds; }
    }

    public int Depth
    {
        get { return m_Depth; }
    }

    public int MaxDepth
    {
        get { return m_MaxDepth; }
    }

    public int GetShapeCount()
    {
        return m_ShapeList.Count;
    }

    public IShape GetShape(int i)
    {
        return m_ShapeList[i];
    }

    public MyQuadTree GetChild(int index)
    {
        if (null != m_Children)
            return m_Children[index];
        return null;
    }

    public void AddShape(IShape shape)
    {
        if (null != m_Children)
        {
            AddShapeToBelong(shape);
            return;
        }

        m_ShapeList.Add(shape);
        if (m_ShapeList.Count > m_SplitThreshold)
        {
            if (m_Depth >= m_MaxDepth)
            {
                //Debug.LogWarning($"split fail: reach maxDepth:{m_MaxDepth}");
                return;
            }
            m_Children = new MyQuadTree[4];
            Split();

            //形状重新分配到新分裂的节点中
            for (int i = 0; i < m_ShapeList.Count; ++i)
            {
                AddShapeToBelong(m_ShapeList[i]);
            }
            m_ShapeList.Clear();
        }
    }

    //节点一分为4, 每个节点占据的大小为原来的1/4
    private void Split()
    {
        int nextDepth = m_Depth + 1;
        float halfWidth = m_Bounds.width * 0.5f;
        float halfHeight = m_Bounds.height * 0.5f;
        float xMid = m_Bounds.xMin + halfWidth;
        float yMid = m_Bounds.yMin + halfHeight;

        var leftBottomRect = new Rect(m_Bounds.xMin, m_Bounds.yMin, halfWidth, halfHeight);
        m_Children[0] = new MyQuadTree(leftBottomRect, m_SplitThreshold, m_MaxDepth);
        m_Children[0].m_Depth = nextDepth;

        var leftTopRect = new Rect(m_Bounds.xMin, yMid, halfWidth, halfHeight);
        m_Children[1] = new MyQuadTree(leftTopRect, m_SplitThreshold, m_MaxDepth);
        m_Children[1].m_Depth = nextDepth;

        var rightTopRect = new Rect(xMid, yMid, halfWidth, halfHeight);
        m_Children[2] = new MyQuadTree(rightTopRect, m_SplitThreshold, m_MaxDepth);
        m_Children[2].m_Depth = nextDepth;

        var rightBottomRect = new Rect(xMid, m_Bounds.yMin, halfWidth, halfHeight);
        m_Children[3] = new MyQuadTree(rightBottomRect, m_SplitThreshold, m_MaxDepth);
        m_Children[3].m_Depth = nextDepth;
    }

    private int[] m_TempIndex = new int[4];
    //获取包围盒所属的节点索引
    private int[] GetAABBBelongIndexes(Rect aabb)
    {
        int index = 0;
        if (aabb.Overlaps(m_Bounds))
        {
            float halfWidth = m_Bounds.width * 0.5f;
            float halfHeight = m_Bounds.height * 0.5f;
            float xMid = m_Bounds.xMin + halfWidth;
            float yMid = m_Bounds.yMin + halfHeight;
            var left = aabb.xMin < xMid;
            var right = aabb.xMax > xMid;
            var bottom = aabb.yMin < yMid;
            var top = aabb.yMax > yMid;

            if (left)
            {
                if (bottom)
                    m_TempIndex[index++] = 0;
                if (top)
                    m_TempIndex[index++] = 1;
            }
            if (right)
            {
                if (top)
                    m_TempIndex[index++] = 2;
                if (bottom)
                    m_TempIndex[index++] = 3;
            }
        }
        for (int i = index; i < 4; ++i)
            m_TempIndex[i] = -1;

        return m_TempIndex;
    }

    //添加到所属的节点, 可以属于多个节点
    private void AddShapeToBelong(IShape shape)
    {
        var shapeBounds = shape.GetBounds();
        var allIndexes = GetAABBBelongIndexes(shapeBounds);
        for (int i = 0; i < allIndexes.Length; ++i)
        {
            int index = allIndexes[i];
            if (-1 == index) break;
            m_Children[index].AddShape(shape);
        }
    }

    private HashSet<IShape> m_TempSet = new HashSet<IShape>();
    //获取与该包围盒发生碰撞的所有Shape
    public List<IShape> Query(Rect bounds)
    {
        m_TempSet.Clear();
        Query(bounds, m_TempSet);
        return new List<IShape>(m_TempSet);
    }

    private void Query(Rect bounds, HashSet<IShape> result)
    {
        //一个Shape可能会添加到多个节点, 用set达到去重效果
        for (int i = 0; i < m_ShapeList.Count; ++i)
            result.Add(m_ShapeList[i]);

        var allIndexes = GetAABBBelongIndexes(bounds);
        if (null != m_Children)
        {
            for (int i = 0; i < allIndexes.Length; ++i)
            {
                int index = allIndexes[i];
                if (-1 == index) break;
                m_Children[index].Query(bounds, result);
            }
        }
    }

    public void Clear()
    {
        m_ShapeList.Clear();

        if (null != m_Children)
        {
            for (int i = 0; i < m_Children.Length; ++i)
                m_Children[i].Clear();
            m_Children = null;
        }
    }

}

using System;
using System.Collections.Generic;
using UnityEditor;
using UnityEngine;

public class MyQuadTreeTest : MonoBehaviour
{
    public Transform m_ShapesRoot;
    public GameObject m_RectTemplate;

    [Range(0, 9)]
    public int m_ApiType = 1;

    public bool m_Benchmark = false; //测试调用多少次时, 执行耗时超过800ms
    [Range(0, 999)]
    public int m_TestShapeCount = 10;

    public int m_CostTime = 0;
    public int m_CheckCount = 0;

    private Vector3 m_PointCubeSize = new Vector3(0.1f, 0.1f, 0.01f);

    private List<OBBRect> m_RectList = new List<OBBRect>();
    private HashSet<OBBRect> m_MouseIntersectRect = new HashSet<OBBRect>(); //SceneView下鼠标下方要标红的矩形
    private MyQuadTree m_Tree;

    void Start()
    {
        m_Tree = new MyQuadTree(new Rect(-10, -10, 20, 20), 4, 4);
        for (int i = 0; i < m_ShapesRoot.childCount; ++i)
        {
            var shape = m_ShapesRoot.GetChild(i).GetComponent<OBBRect>();
            if (null == shape || !shape.isActiveAndEnabled) continue;
            m_RectList.Add(shape);
            m_Tree.AddShape(shape);
        }
    }

    void Update()
    {
        for (int i = 0; i < 10; ++i)
        {
            if (m_RectList.Count >= m_TestShapeCount) break;
            CreateRect();
        }

        switch (m_ApiType)
        {
        case 1:
            TreeCheckIntersect();
            break;

        }

        CheckTimeCost();
    }

    void TreeCheckIntersect()
    {
        m_Tree.Clear();
        for (int i = 0; i < m_RectList.Count; ++i)
        {
            var shape = m_RectList[i];
            if (m_MouseIntersectRect.Contains(shape))
                shape.m_GizmosColor = Color.red;
            else
                shape.m_GizmosColor = Color.white;
            m_Tree.AddShape(shape);
        }

        var t = DateTime.Now;
        m_CheckCount = 0;
        for (int i = 0; i < m_RectList.Count; ++i)
        {
            var shape = m_RectList[i];

            var intersectList = m_Tree.Query(shape.GetBounds());
            if (null != intersectList)
            {
                for (int j = 0; j < intersectList.Count; ++j)
                {
                    var shape_2 = intersectList[j] as OBBRect;
                    if (null == shape_2 || shape == shape_2) continue;
                    m_CheckCount++;

                    if (Shape2DHelper.IsPolygonIntersect(shape.GetVerts(), shape_2.GetVerts()))
                    {
                        shape.m_GizmosColor = Color.red;
                        shape_2.m_GizmosColor = Color.red;
                    }
                }
            }
        }

        m_CostTime = (int)(DateTime.Now - t).TotalMilliseconds;
    }

    protected void CheckTimeCost()
    {
        if (m_CostTime >= 200) //防止卡住
        {
            for (int i = 1; i <= 10; ++i)
            {
                var childIndex = m_ShapesRoot.childCount - i;
                if (childIndex < 0) break;
                m_ShapesRoot.GetChild(childIndex).gameObject.SetActive(false);
            }
        }

        if (m_Benchmark)
        {
            if (m_CostTime >= 100)
            {
                if (m_ApiType >= 2)
                {
                    m_ApiType = 1;
                    m_Benchmark = false;
                }
                else
                    m_ApiType++;
            }
            else
            {
                if (m_RectList.Count >= m_TestShapeCount)
                {
                    m_TestShapeCount++;
                }
            }
        }
    }

    private void CreateRect()
    {
        var go = GameObject.Instantiate(m_RectTemplate, m_ShapesRoot, false);
        go.SetActive(true);
        go.name = $"Rect ({m_RectList.Count})";
        var trans = go.transform;
        trans.position = new Vector2(UnityEngine.Random.Range(-10, 10), UnityEngine.Random.Range(-10, 10));
        trans.localEulerAngles = new Vector3(0, 0, UnityEngine.Random.Range(-30, 30));

        var shape = go.GetComponent<OBBRect>();
        shape.m_Width = UnityEngine.Random.Range(0.1f, 3f);
        shape.m_Height = shape.m_Width * UnityEngine.Random.Range(0.6f, 0.8f);
        m_RectList.Add(shape);
        m_Tree.AddShape(shape);
    }

    private void CreateRect2(Vector3 pos)
    {
        var go = GameObject.Instantiate(m_RectTemplate, m_ShapesRoot, false);
        go.SetActive(true);
        go.name = $"Rect ({m_RectList.Count})";
        var trans = go.transform;
        trans.position = pos;
        trans.localEulerAngles = new Vector3(0, 0, UnityEngine.Random.Range(-30, 30));

        var shape = go.GetComponent<OBBRect>();
        shape.m_Width = 2;
        shape.m_Height = 1.5f;
        m_RectList.Add(shape);
        m_Tree.AddShape(shape);
    }

#if UNITY_EDITOR

    private Stack<MyQuadTree> m_DrawStack = new Stack<MyQuadTree>();
    private void OnDrawGizmos()
    {
        DrawGizmos_Tree();
    }

    private void DrawGizmos_Tree()
    {
        if (null != m_Tree)
        {
            var curEvent = Event.current;
            var cam = Camera.current;
            if (null != cam)
            {
                Vector3 screenPos = HandleUtility.GUIPointToScreenPixelCoordinate(curEvent.mousePosition); //转换成像素单位以及左下角(0, 0)
                screenPos.z = -cam.transform.position.z;
                var worldPos = cam.ScreenToWorldPoint(screenPos);
                //Debug.Log($"mousePos:{curEvent.mousePosition} -> screenPos:{screenPos}, worldPos:{worldPos}");

                float handleSizeFactor = HandleUtility.GetHandleSize(worldPos);
                var pointCubeSize = m_PointCubeSize * handleSizeFactor;
                Gizmos.DrawCube((Vector2)worldPos, pointCubeSize);

                var mousePosBounds = new Rect(worldPos.x - m_PointCubeSize.x * 0.5f, worldPos.y - m_PointCubeSize.y * 0.5f, m_PointCubeSize.x, m_PointCubeSize.y);
                var intersectList = m_Tree.Query(mousePosBounds);
                m_MouseIntersectRect.Clear();
                if (null != intersectList)
                {
                    for (int i = 0; i < intersectList.Count; ++i)
                    {
                        var shape = intersectList[i] as OBBRect;
                        m_MouseIntersectRect.Add(shape);
                    }
                }
                if (curEvent.type == EventType.MouseUp) //在鼠标位置添加一个矩形
                {
                    CreateRect2(worldPos);
                    curEvent.Use();
                }
            }

            var node = m_Tree;
            m_DrawStack.Push(node);
            do
            {
                node = m_DrawStack.Pop();
                var c = Color.HSVToRGB(node.Depth / (float)node.MaxDepth, 1, 1);
                //c.a = 0.5f;
                Gizmos.color = c;
                DrawNodeBounds(node);

                for (int i = 0; i < 4; ++i)
                {
                    var child = node.GetChild(i);
                    if (null != child)
                        m_DrawStack.Push(child);
                }
            } while (m_DrawStack.Count > 0);

            Gizmos.color = Color.white;
        }
    }

    private void DrawNodeBounds(MyQuadTree node)
    {
        var bounds = node.Bounds;
        //显示的时候, 左下角padding缩小一点, 不然线叠在一起看不清
        var rootBounds = m_Tree.Bounds;
        float x = node.Depth * rootBounds.width * 0.003f;
        float y = node.Depth * rootBounds.height * 0.003f;
        bounds.xMin += x;
        bounds.yMin += y;

        var min = bounds.min;
        var max = bounds.max;
        var leftTop = new Vector2(min.x, max.y);
        var rightBottom = new Vector2(max.x, min.y);
        Gizmos.DrawLine(min, leftTop);
        Gizmos.DrawLine(leftTop, max);
        Gizmos.DrawLine(max, rightBottom);
        Gizmos.DrawLine(rightBottom, min);
    }

#endif

}