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