00001
00002
00003
00004
00005
00006
00007
00008
00009
00010
00011
00012
00013
00014
00015
00016
00017
00018
00019
00020
00021
00022
00023 #ifndef LUX_KDTREE_H
00024 #define LUX_KDTREE_H
00025
00026 #include "lux.h"
00027 #include "memory.h"
00028
00029
00030 namespace lux
00031 {
00032
00033 struct KdNode {
00034 void init(float p, u_int a) {
00035 splitPos = p;
00036 splitAxis = a;
00037
00038 rightChild = 0;
00039 rightChild = ~rightChild;
00040 hasLeftChild = 0;
00041 }
00042 void initLeaf() {
00043 splitAxis = 3;
00044
00045 rightChild = 0;
00046 rightChild = ~rightChild;
00047 hasLeftChild = 0;
00048 }
00049
00050 float splitPos;
00051 u_int splitAxis:2;
00052 u_int hasLeftChild:1;
00053 u_int rightChild:29;
00054 };
00055 template <class NodeData, class LookupProc> class KdTree {
00056 public:
00057
00058 KdTree(const vector<NodeData> &data);
00059 ~KdTree() {
00060 FreeAligned(nodes);
00061 delete[] nodeData;
00062 }
00063 void recursiveBuild(u_int nodeNum, int start, int end,
00064 vector<const NodeData *> &buildNodes);
00065 void Lookup(const Point &p, const LookupProc &process,
00066 float &maxDistSquared) const;
00067 NodeData *getNodeData() { return nodeData; }
00068
00069 private:
00070
00071 void privateLookup(u_int nodeNum, const Point &p,
00072 const LookupProc &process, float &maxDistSquared) const;
00073
00074 KdNode *nodes;
00075 NodeData *nodeData;
00076 u_int nNodes, nextFreeNode;
00077 };
00078 template<class NodeData> struct CompareNode {
00079 CompareNode(int a) { axis = a; }
00080 int axis;
00081 bool operator()(const NodeData *d1,
00082 const NodeData *d2) const {
00083 return d1->p[axis] == d2->p[axis] ? (d1 < d2) :
00084 d1->p[axis] < d2->p[axis];
00085 }
00086 };
00087
00088 template <class NodeData, class LookupProc>
00089 KdTree<NodeData,
00090 LookupProc>::KdTree(const vector<NodeData> &d) {
00091 nNodes = d.size();
00092 nextFreeNode = 1;
00093 nodes = AllocAligned<KdNode>(nNodes);
00094 nodeData = new NodeData[nNodes];
00095 vector<const NodeData *> buildNodes;
00096 for (u_int i = 0; i < nNodes; ++i)
00097 buildNodes.push_back(&d[i]);
00098
00099 recursiveBuild(0, 0, nNodes, buildNodes);
00100 }
00101 template <class NodeData, class LookupProc> void
00102 KdTree<NodeData, LookupProc>::recursiveBuild(u_int nodeNum,
00103 int start, int end,
00104 vector<const NodeData *> &buildNodes) {
00105
00106 if (start + 1 == end) {
00107 nodes[nodeNum].initLeaf();
00108 nodeData[nodeNum] = *buildNodes[start];
00109 return;
00110 }
00111
00112
00113 BBox bound;
00114 for (int i = start; i < end; ++i)
00115 bound = Union(bound, buildNodes[i]->p);
00116 int splitAxis = bound.MaximumExtent();
00117 int splitPos = (start+end)/2;
00118 std::nth_element(&buildNodes[start], &buildNodes[splitPos],
00119 &buildNodes[end-1], CompareNode<NodeData>(splitAxis));
00120
00121
00122 nodes[nodeNum].init(buildNodes[splitPos]->p[splitAxis],
00123 splitAxis);
00124 nodeData[nodeNum] = *buildNodes[splitPos];
00125 if (start < splitPos) {
00126 nodes[nodeNum].hasLeftChild = 1;
00127 u_int childNum = nextFreeNode++;
00128 recursiveBuild(childNum, start, splitPos, buildNodes);
00129 }
00130 if (splitPos+1 < end) {
00131 nodes[nodeNum].rightChild = nextFreeNode++;
00132 recursiveBuild(nodes[nodeNum].rightChild, splitPos+1,
00133 end, buildNodes);
00134 }
00135 }
00136 template <class NodeData, class LookupProc> void
00137 KdTree<NodeData, LookupProc>::Lookup(const Point &p,
00138 const LookupProc &proc,
00139 float &maxDistSquared) const {
00140 privateLookup(0, p, proc, maxDistSquared);
00141 }
00142 template <class NodeData, class LookupProc> void
00143 KdTree<NodeData, LookupProc>::privateLookup(u_int nodeNum,
00144 const Point &p, const LookupProc &process,
00145 float &maxDistSquared) const {
00146 KdNode *node = &nodes[nodeNum];
00147
00148 int axis = node->splitAxis;
00149 if (axis != 3) {
00150 float dist = p[axis] - node->splitPos;
00151 float dist2 = dist * dist;
00152 if (p[axis] <= node->splitPos) {
00153 if (node->hasLeftChild)
00154 privateLookup(nodeNum+1, p,
00155 process, maxDistSquared);
00156 if (dist2 < maxDistSquared &&
00157 node->rightChild < nNodes)
00158 privateLookup(node->rightChild,
00159 p,
00160 process,
00161 maxDistSquared);
00162 }
00163 else {
00164 if (node->rightChild < nNodes)
00165 privateLookup(node->rightChild,
00166 p,
00167 process,
00168 maxDistSquared);
00169 if (dist2 < maxDistSquared && node->hasLeftChild)
00170 privateLookup(nodeNum+1,
00171 p,
00172 process,
00173 maxDistSquared);
00174 }
00175 }
00176
00177 float dist2 = DistanceSquared(nodeData[nodeNum].p, p);
00178 if (dist2 < maxDistSquared)
00179 process(nodeData[nodeNum], dist2, maxDistSquared);
00180 }
00181
00182 }
00183
00184 #endif // LUX_KDTREE_H