愿知识带给你我勇气,用那温暖的光,驱散冰冷的黑夜。

MENU

【代码札记】Kotlin实现四叉树寻找区域内的点

May 19, 2022 • Read: 428 • 瞎折腾

本文记述如何使用Kotlin实现一个PR四叉树(点-区域四叉树),并用该四叉树实现搜索给定区域中的点。

序言

四叉树和二叉树类似,都是树状结构。二叉树有两个子树,而四叉树有四个,这使得四叉树能够方便的对应上二维平面:将一个区域分割为四块,每块刚好对应一个子树。根据划分方式和存储的值的类型,四叉树有很多变体,例如区域四叉树(永远平均分割)、点四叉树(每次以某个点为中心进行区域分割)、点-区域四叉树(与区域二叉树类似,但叶节点存储其中的所有节点)、边四叉树等,关于四叉树的详细介绍这里就不多赘述了,想了解更多可以移步维基百科

本文要实现的是点-区域四叉树。它在划分方式上与区域二叉树相同:永远在正中心将区域四等分。而它与区域二叉树不同的是,区域二叉树每个节点仅存储一个值,表示该区域的平均量;而点-区域二叉树在节点内存储所有包含在该片区域的点。因此很适合用来搜索给定区域中包含的所有点。

关于使用四叉树的好处这里就不多说了,因为很简单:如果不使用四叉树,要寻找位于指定区域内的点就需要顺序遍历所有点,而使用四叉树则可以直接排除掉部分点,因为这些子树所代表的区域与待查找区域没有交集,自然也没有符合条件的点。

模型定义

在实现四叉树之前,我们需要一些数据结构。处理二位平面就离不开坐标,为了保持通用性,我们将核心的坐标定义为一个接口XY

interface XY {
    val x: Int
    val y: Int
}

在本文中我们的操作对象仅限于点,因此实现一个点:

data class Point(
    override val x: Int,
    override val y: Int
) : XY

然后我们还得想办法描述一个区域,关于区域的描述有很多种,比如:

  • 给定中心点,分别给出半长和半宽
  • 给定一个顶点,分别给出长和宽
  • 给定两个顶点

本文选用第三种方法,直接给定两个顶点:

data class BoundingBox(
    val topLeftX: Int,
    val topLeftY: Int,
    val bottomRightX: Int,
    val bottomRightY: Int,
)

而基于面向对象的方法,这个BoundingBox应当还有一些成员方法,例如判断一个点是否在其内部,是否与另一个区域相交,以及将自己划分为四个子区域:

data class BoundingBox(
    val topLeftX: Int,
    val topLeftY: Int,
    val bottomRightX: Int,
    val bottomRightY: Int,
) {
    /**
     * Test if a point is in the range of the box.
     * Return true if the point is in the box.
     *
     * The top left corner is (0,0), the bottom right corner is (x,y).
     * Also, the top left edges are included, but bottom right edges are excluded.
     * */
    fun containsPoint(point: XY): Boolean {
        return point.x in (topLeftX until bottomRightX) && point.y in (topLeftY until bottomRightY)
    }

    /**
     * Test if the target box overlaps with this box.
     *
     * By "overlap", there must be shared area, not shared edges.
     * So just share the same edges doesn't mean overlap.
     * */
    fun overlap(target: BoundingBox): Boolean {
        if (target.bottomRightX < this.topLeftX) {
            // target on the left
            return false
        }
        if (target.topLeftX > this.bottomRightX) {
            // target on the right
            return false
        }
        if (target.bottomRightY < this.topLeftY) {
            // target on the top
            return false
        }
        if (target.topLeftY > this.bottomRightY) {
            // target on the bottom
            return false
        }

        return true
    }

    /**
     * Split this box and get 4 sub boxes.
     * Index:
     * + 0: top left
     * + 1: top right
     * + 2: bottom left
     * + 3: bottom right
     * */
    fun getSubRegion(index: Int): BoundingBox {
        val newBoxWidth = (this.bottomRightX - this.topLeftX) / 2
        val newBoxHeight = (this.bottomRightY - this.topLeftY) / 2

        return when (index) {
            0 -> BoundingBox(
                this.topLeftX, this.topLeftY,
                this.topLeftX + newBoxWidth, this.topLeftY + newBoxHeight
            )
            1 -> BoundingBox(
                this.topLeftX + newBoxWidth, this.topLeftY,
                this.bottomRightX, this.topLeftY + newBoxHeight
            )
            2 -> BoundingBox(
                this.topLeftX, this.topLeftY + newBoxHeight,
                this.topLeftX + newBoxWidth, this.bottomRightY
            )
            3 -> BoundingBox(
                this.topLeftX + newBoxWidth, this.topLeftY + newBoxHeight,
                this.bottomRightX, this.bottomRightY
            )
            else -> throw IllegalArgumentException("The index must in range [0,3]")
        }
    }

其中containsPoint判断区域中是否包含某个点,判断坐标就是了。而overlap则判断是否与另一个边界有重叠:排除四种不可能重叠的情况,那剩下的就是重叠了。最后getSubRegion根据序号0到3,分别对应四个子区域,这样设计是为了初始化四叉树节点设计的(直接遍历0到3,而不是设定4个枚举类型)。

铺垫好了数据结构,下面我们说一说四叉树怎么写。

四叉树实现

基本框架

首先是树的基本内容:

class QuadTree(
    /**
     * The region of this quad tree representing.
     * */
    val region: BoundingBox,

    /**
     * The max amount of elements can contain before this quad tree split.
     * */
    private val capacity: Int,
) {
    private val points: MutableList<XY> = ArrayList(capacity)

    /**
     * Subtrees, order: [topLeft, topRight, bottomLeft, bottomRight]
     * */
    private var _subTrees: Array<QuadTree>? = null

    val subTrees: List<QuadTree>
        get() = _subTrees?.toList() ?: emptyList()

    /**
     * Insert a point into this node.
     * */
    fun insert(point: XY): Boolean {
        TODO()
    }

    /**
     * Return all points in the given range
     * */
    fun searchRange(target: BoundingBox): List<XY> {
        TODO()
    }

}

目前这个类包含了一些基本的成员,例如当前节点表示的区域、节点容量(在点过多时分裂节点)、点列表、子树列表。这里子树列表用的数组表示,最初为null,仅在需要时才分裂节点。但是当外部请求的时候,我们希望给的是不可变列表,从而避免调用者拿着返回值乱搞,所以将内部的成员起名为_subTrees,而subTrees则对应的是其不可变列表。

后面还定义了两个基本操作:插入(构建树)和查询。其中查询比较简单,我们先说说查询。

查询

查询功能将根据给定的区域,搜索以本节点为树根的四叉树中,所有位于给定区域中的点。这个还是很好实现的:

    fun searchRange(target: BoundingBox): List<XY> {
        if (!this.region.overlap(target)) {
            // no overlaps, no points
            return emptyList()
        }
        return this._subTrees
            ?.flatMap { it.searchRange(target) }
            ?: this.points.filter { target.containsPoint(it) }
    }

首先判断给定区域与当前区域是否有重叠——如果没有,那肯定查不到节点,直接返回空列表,这就相当于剪枝。之后根据当前节点的状态,有两种情况:

  • 若子树为空,则本节点为叶子节点,存储了点信息
  • 若子树不为空,则本节点已经分裂,不包含任何点信息

若子树不为空,则查找四棵子树,然后整理成不可变列表返回;否则搜索当前节点内的点。

插入

说完了查询,再说说插入。插入节点也可以分为两种情况:

  • 若当前节点有子树,则直接插入子树节点
  • 若当前节点没有子树

    • 若节点未满,则存入点列表
    • 若节点已满,分裂节点并插入子节点

其中插入子节点可以单独封装成一个方法:

    private fun insertToSubRegions(point: XY): Boolean {
        return this._subTrees?.any { it.insert(point) } ?: false
    }

若任意一个子树接受了点,则返回true,否则返回false。若没有子树,也返回false。这里利用了一个技巧来简化写法,即正常情况下,有且只有一个子树能够接受这个点(需要由当前节点保证该点在当前树枝的范围内),因此子树不为空时any总为true(有且只有一个子树插入成功,返回true),而我们想要的是它在判断true时尝试插入点,这样我们既插入了点,又得到了返回值。

之后即可按照上面的思路编写插入方法了:

    fun insert(point: XY): Boolean {
        if (!this.region.containsPoint(point)) {
            // the point is not in this region, reject it
            return false
        }
        if (insertToSubRegions(point)) {
            return true
        }
        check(_subTrees == null) { "All 4 sub nodes rejected the point" }
        return if (this.points.size < capacity) {
            // now we can just insert and return
            this.points.add(point)
            true
        } else {
            // we need to split/subdivide this node into 4 sub nodes
            _subTrees = Array(4) { QuadTree(region.getSubRegion(it), capacity) }
            // empty the points and insert them in to sub nodes
            points.forEach { check(insertToSubRegions(it)) { "Father node accepted, but all sub nodes rejected" } }
            points.clear()
            // try sub nodes
            insertToSubRegions(point)
        }
    }

在插入前检查范围,如果待插入的点根本不在当前树枝的覆盖范围内,则直接返回false。然后尝试插入子树:插入成功则返回true,插入失败则有两种情况,要么子树为空,要么全部子树拒绝了这个点。后面的check要求子树为空,否则直接抛出IllegalStateException,因为四棵子树必须有且只有一棵接受这个点。检查通过之后就可以确认子树是空的,因此可以先看节点容量,没满则插入点列表,返回true;满了则分裂节点:首先产生4棵子树,然后将现有的点插入子树并清空,最后将新的点插入子树。将现有点插入子树时要求必须插入成功,否则就会出现树根接受了点,而子树拒绝了点的情况。

在实际运行时,这些check可能会拖慢运行速度,但在开发过程中,验证这些条件有助于排错。否则程序无声无息的挂掉了,你都不知道从何处入手。

搜索可视化

写完了树,当然要用用看。既然是平面数据,自然离不开可视化。这里我们随机初始化一些点,然后随机给一个搜索区域,通过四叉树搜索区域中的点,并赋予不同颜色画出来。

package info.skyblond.quadtree

import java.awt.BasicStroke
import java.awt.Color
import java.awt.Graphics2D
import java.awt.image.BufferedImage
import java.awt.RenderingHints
import java.awt.geom.Ellipse2D
import java.io.File
import javax.imageio.ImageIO
import kotlin.random.Random

object QuadTreeDemo {
    private const val width = 360
    private const val height = 240

    @JvmStatic
    fun main(args: Array<String>) {
        val boundary = BoundingBox(0, 0, width, height)
        val quadTree = QuadTree(boundary, 4)

        val dots = (width.toDouble() * height * 0.02).toLong()

        for (i in 1..dots) {
            val p = Point(Random.nextInt(width), Random.nextInt(height))
            quadTree.insert(p)
        }

        // 1 extra for edge stroke
        val image = BufferedImage(width + 1, height + 1, BufferedImage.TYPE_INT_RGB)
        val graphics = image.createGraphics()
        graphics.background = Color.BLACK
        graphics.color = Color.WHITE

        // draw quad tree tiles
        graphics.stroke = BasicStroke(1.0F)
        drawQuadTreeTile(quadTree, graphics)

        // draw all points
        graphics.color = Color.GREEN
        quadTree.searchRange(boundary)
            .forEach {
                fillCircle(it.x, it.y, 1.0, graphics)
            }

        // search region
        val searchWidth = Random.nextInt(width / 2)
        val searchX = Random.nextInt(searchWidth)
        val searchHeight = Random.nextInt(width / 2)
        val searchY = Random.nextInt(searchHeight)
        val searchRegion = BoundingBox(searchX, searchY, searchX + searchWidth, searchY + searchHeight)
        graphics.color = Color.YELLOW
        drawRegion(searchRegion, graphics)

        // draw searched points
        graphics.color = Color.RED
        quadTree.searchRange(searchRegion)
            .forEach { fillCircle(it.x, it.y, 1.0, graphics) }

        // done
        graphics.dispose()
        ImageIO.write(image, "png", File("quadtree_search.png"))
    }

    private fun drawQuadTreeTile(quadTree: QuadTree, graphics: Graphics2D) {
        drawRegion(quadTree.region, graphics)
        quadTree.subTrees.forEach { drawQuadTreeTile(it, graphics) }
    }

    private fun drawRegion(region: BoundingBox, graphics: Graphics2D) {
        val regionWidth = region.bottomRightX - region.topLeftX
        val regionHeight = region.bottomRightY - region.topLeftY
        graphics.drawRect(region.topLeftX, region.topLeftY, regionWidth, regionHeight)
    }
    
    fun fillCircle(centerX: Int, centerY: Int, radius: Double, graphics: Graphics2D) {
        // save the old value
        val oldValue = graphics.getRenderingHint(RenderingHints.KEY_ANTIALIASING)

       graphics.setRenderingHint(RenderingHints.KEY_ANTIALIASING, RenderingHints.VALUE_ANTIALIAS_ON)
        val shape = Ellipse2D.Double(centerX - radius, centerY - radius, 2 * radius, 2 * radius)
        graphics.fill(shape)

        // recover the old value
        graphics.setRenderingHint(RenderingHints.KEY_ANTIALIASING, oldValue)
    }
}

一次运行的结果如下所示:

test.png

其中白框为四叉树各节点的覆盖范围,绿色点为普通点,黄色框为搜索框,红色点为四叉树搜索到的点。可见效果还是挺好的。

至于四叉树的性能嘛,如果只是寻找点的话,实际上它是比不过遍历列表的。因为遍历列表是循环,而查找树则是递归,除了保存栈帧等较慢的操作之外,还可能会出现栈内存溢出(点太多了,递归调用直接爆栈)。下一篇文章将介绍如何使用四叉树进行圆圈的碰撞检测,在这个过程中需要对每个圆搜索其附近的圆,如果使用列表的话,N个圆就是N^2次搜索,而使用树则近似为Nlog4 N,这时四叉树的优势就体现出来了。

四叉树完整代码

下面是四叉树的完整代码:

package info.skyblond.quadtree


class QuadTree(
    /**
     * The region of this quad tree representing.
     * */
    val region: BoundingBox,

    /**
     * The max amount of elements can contain before this quad tree split.
     * */
    private val capacity: Int,
) {
    private val points: MutableList<XY> = ArrayList(capacity)

    /**
     * Subtrees, order: [topLeft, topRight, bottomLeft, bottomRight]
     * */
    private var _subTrees: Array<QuadTree>? = null

    val subTrees: List<QuadTree>
        get() = _subTrees?.toList() ?: emptyList()

    /**
     * Insert a point into this node.
     * */
    fun insert(point: XY): Boolean {
        if (!this.region.containsPoint(point)) {
            // the point is not in this region, reject it
            return false
        }
        if (insertToSubRegions(point)) {
            return true
        }
        check(_subTrees == null) { "All 4 sub nodes rejected the point" }
        return if (this.points.size < capacity) {
            // now we can just insert and return
            this.points.add(point)
            true
        } else {
            // we need to split/subdivide this node into 4 sub nodes
            _subTrees = Array(4) { QuadTree(region.getSubRegion(it), capacity) }
            // empty the points and insert them in to sub nodes
            points.forEach { check(insertToSubRegions(it)) { "Father node accepted, but all sub nodes rejected" } }
            points.clear()
            // try sub nodes
            insertToSubRegions(point)
        }
    }

    /**
     * Insert to subregions
     * */
    private fun insertToSubRegions(point: XY): Boolean {
        return this._subTrees?.any { it.insert(point) } ?: false
    }

    /**
     * Return all points in the given range
     * */
    fun searchRange(target: BoundingBox): List<XY> {
        if (!this.region.overlap(target)) {
            // no overlaps, no points
            return emptyList()
        }
        return this._subTrees
            ?.flatMap { it.searchRange(target) }
            ?: this.points.filter { target.containsPoint(it) }
    }
}

-全文完-


知识共享许可协议
【代码札记】Kotlin实现四叉树寻找区域内的点天空 Blond 采用 知识共享 署名 - 非商业性使用 - 相同方式共享 4.0 国际 许可协议进行许可。
本许可协议授权之外的使用权限可以从 https://www.skyblond.info/about.html 处获得。

Archives QR Code
QR Code for this page
Tipping QR Code