KD树 (k-Dimension tree)

1、简介

KD树是对k维空间进行分割，引用于k维空间的检索的一种数据结构

利用k维空间上的点进行分割，组成二叉树，就是KD树，一般将其分割成平衡二叉树。

统计学习方法一书中使用的KD树的例子：

形成的KD树为：

当然，节点中的每个点表示一个k-1维的超平面（2维空间里面就变成直线了），每个节点还需要一个方向来表示该平面的法向量。

分割方向的选择一般有两种，一是xyz轮流，另一种是取该方向方差最大（分布最分散）的。

2、实现

尝试用scala写一个KD树的实现。

首先是定义：

case class KDTree(left:KDTree,right:KDTree,point:Array[Int],direction:Int){

  def isLeaf = left==null && right==null

  def preOrder: Unit={

    println(point.mkString(","))

    if(left!=null) left.preOrder

    if(right!=null) right.preOrder

  }

}

顺便带了个前序遍历，方便查看测试结果

建树，为了方便，这里是递归版本，采用轮流的方式选取法向量：

def create(ps: Seq[Array[Int]], depth:Int):KDTree={

    if(ps.isEmpty) return null

    val numPoints = ps.size

    val direction = depth % ps.head.length

    val med = ps.sortBy(_(direction)).apply(numPoints/2)

    val leftSeq = ps.filter(_(direction)<med(direction))

    val rightSeq = ps.filter(_(direction)>med(direction))

    new KDTree(create(leftSeq,depth+1), create(rightSeq,depth+1), med, direction)

  }

最后，用1中的数据做一个测试

val kdtree = KDTree.create(List(Array(2,3),Array(5,4),Array(4,7),Array(7,2),Array(9,6),Array(8,1)), 0)

  kdtree.preOrder

结果如下：

3、变化

有时候点在不同维度分布很不均匀，极端情况，比如全部分布在一条直线上，直接这样就会得到不好的结果，这时候就可以取方差最大的一维进行分割。

附录中有一个非递归、取方差的scala代码。

用(1,1)(1,2)...(1,7)共7个点做测试，由于每次都是y方向方差最大，所以结果应该是排序树的样子，如下：

附录，非递归，取方差最大方向为轴的kd树的构造代码

case class Point(coo:Array[Int], label:Int)

 

case class KDTree(var left:KDTree,var right:KDTree,var points:Seq[Point],var direction:Int){

  def this() = this(null,null,null,-1)

  def isLeaf = left==null && right==null

  def preOrder: Unit={

    var list = List[KDTree]()

    list ::= this

    var tmp:KDTree = null

    while(!list.isEmpty){

      tmp = list.head

      list = list.tail

      println('('+tmp.points.map(_.label).mkString(",")+')')

      if(tmp.right!=null) list ::= tmp.right

      if(tmp.left!=null) list ::= tmp.left

    }

  }

}

 

object KDTree{

  def create(ps: Seq[Point]):KDTree={

    val dimension = ps.head.coo.length

    val root = new KDTree

    var list = List[(KDTree,Seq[Point])]((root, ps))

    while(!list.isEmpty){

      val (node, spoints) = list.head

      list = list.tail

      val n = spoints.size

      var max = 0.0

      for(i<-0 until dimension){

        val (t1,t2) = spoints.map(_.coo(i)).map(x=>(x,x*x)).reduce((a,b)=>(a._1+b._1, a._2+b._2))

        val tmp = (t2-t1*t1.toDouble/n)/n

        if(max<=tmp){

          max = tmp

          node.direction = i

        }

      }

      val sortedPoints = spoints.sortBy(_.coo(node.direction))

      val med = sortedPoints(n/2).coo(node.direction)

      val i1 = sortedPoints.indexWhere(_.coo(node.direction)==med)

      val i2 = sortedPoints.indexWhere(_.coo(node.direction)!=med, i1+1)

      node.points = sortedPoints.slice(i1, if(i2>0)i2 else n)

      if(i1!=0){

        val left = new KDTree

        node.left = left

        list = (left, sortedPoints.slice(0, i1))::list

      }

      if(i2<n && i2>0){

        val right = new KDTree

        node.right = right

        list = (right, sortedPoints.slice(i2, n))::list

      }

    }

    root

  }

}