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) } |
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 |
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 } } |