[ActionScript 3.0] k-d木

平面上に複数の点が存在する場合、k-d木に格納する事によって、特定領域（矩形）に含まれる点の集合を高速に探索できる。
下のデモでは3000個の点の中から特定の矩形に含まれる点を探索している。

Main.as

package
{
    import flash.display.*;
    import flash.geom.*;
    import flash.events.*;
    import flash.text.*;
    import flash.utils.*;

    [SWF(width="500", height="500", backgroundColor="#ffffff")]
        public class Main extends MovieClip
        {
            public var objects:Array;
            public var objectsInRect:Array;
            private var kd:KDTree;

            private var vx:Number;
            private var vy:Number;

            private var rect:Rectangle;

            private var tf:TextField;

            public function Main()
            {
                tf = new TextField();
                addChild(tf);

                kd = new KDTree();
                rect = new Rectangle(0,0,50,50);
                vx = Math.random()*50;
                vy = Math.random()*50;

                // create some simple objects
                createObjects(3000);

                objectsInRect=[];

                var t:Timer = new Timer(300);
                t.addEventListener(TimerEvent.TIMER, loop);
                t.start();

            }

            private function loop(event:Event):void
            {
                for each(var c:Circle in objectsInRect){
                    c.color = 0xff;
                }

                rect.x += vx;
                rect.y += vy;
                if(rect.x <= 0 || rect.x + rect.width >= stage.stageWidth)
                    vx *= -1;
                if(rect.y <= 0 || rect.y + rect.height >= stage.stageHeight)
                    vy *= -1;

                graphics.clear();
                graphics.lineStyle(1,0,0.5);
                graphics.drawRect(rect.x, rect.y, rect.width, rect.height);

                objectsInRect = [];
                kd.searchPoint(rect, objectsInRect);
                tf.text = kd.searchCount.toString();
                for each(c in objectsInRect){
                    c.color = 0xff0000;
                }
            }

            private function createObjects(count:Number):void
            {
                objects = [];

                for (var i:Number = 0; i < count; i++)
                {
                    newCircle();
                }

                createKDTree(objects, true);
            }

            private function createKDTree(a:Array, isX:Boolean):void{
                var lhs:Array = [];
                var rhs:Array = [];
                var middle:Circle = divideAtMiddlePoint(a, isX, lhs, rhs);
                trace(middle.x + ":" + middle.y);
                kd.insertPoint(middle, true);

                if(lhs.length > 0)
                    createKDTree(lhs, !isX);
                if(rhs.length > 0)
                    createKDTree(rhs, !isX);
            }

            private function divideAtMiddlePoint(original:Array, isX:Boolean, lhs:Array, rhs:Array):Circle{
                var middle:Circle;
                if(isX)
                    original.sort(sortByX);
                else
                    original.sort(sortByY);

                if(original.length >= 3){
                    var index:int = Math.floor(original.length/2);
                    middle = original[index];
                    var i:int;
                    for(i = 0 ; i < index ; i++){
                        lhs.push(original[i]);
                    }
                    for(i = index + 1 ; i < original.length ; i++){
                        rhs.push(original[i]);
                    }
                }else{
                    middle = original[0];
                    for(i = 1 ; i < original.length ; i++){
                        rhs.push(original[i]);
                    }
                }
                return middle;
            }

            private function sortByX(lhs:Circle, rhs:Circle):int{
                return (lhs.x - rhs.x);
            }

            private function sortByY(lhs:Circle, rhs:Circle):int{
                return (lhs.y - rhs.y);
            }

            private function newCircle():void{
                var radius:Number = 2;
                var x:Number = rand(radius, stage.stageWidth - radius);
                var y:Number = rand(radius, stage.stageHeight - radius);
                var circle:Circle = new Circle(x, y, radius);
                addChild(circle);
                circle.x = x;
                circle.y = y;
                circle.color = 0xff;
                objects.push(circle);
            }

            private function rand(min:int, max:int):int
            {
                return min + Math.floor(Math.random() * (max - min + 1));
            }
        }
}

KDTree.as

// 左上を原点とした座標で大小比較していることに注意
package{
    import flash.geom.*;
    public class KDTree{
        private var root:Node = null;
        public var searchCount:int = 0;
        public function KDTree(){
        }

        public function insertPoint(p:Circle, isEven:Boolean):void{
            root = insert(root, isEven, p);
        }

        private function insert(node:Node, isEven:Boolean, p:Circle):Node{
            if(node == null)
                return new Node(p);
            if((isEven && p.x < node.p.x) || (!isEven && p.y < node.p.y)){
                trace("try left:" + isEven + ":(" + node.p.x + "," + node.p.y + ")");
                node.left = insert(node.left, !isEven, p);
            }
            else{
                trace("try right:" + isEven + ":(" + node.p.x + "," + node.p.y + ")");
                node.right = insert(node.right, !isEven, p);
            }

            return node;
        }

        public function searchPoint(r:Rectangle, result:Array):void{
            searchCount = 0;
            search(root, true, r, result);
        }

        private function search(node:Node, isEven:Boolean, r:Rectangle, result:Array):void{
            if (node == null) return;
            trace(node.p.x + ":" + node.p.y);
            searchCount++;
            if(r.left <= node.p.x && node.p.x <= r.right &&
                    r.top <= node.p.y && node.p.y <= r.bottom)
                result.push(node.p);
            if((isEven && node.p.x > r.left) || (!isEven && node.p.y > r.top)){
                trace("search left");
                search(node.left, !isEven, r, result);
            }
            if((isEven && node.p.x < r.right) || (!isEven && node.p.y < r.bottom)){
                trace("search right");
                search(node.right, !isEven, r, result);
            }
        }

    }

}

internal class Node{
    public var p:Circle;
    public var left:Node = null;
    public var right:Node = null;
    public function Node(p:Circle){
        this.p = p;
    }
}