KdTree.hpp

Go to the documentation of this file.

 
namespace GfxTL
{
    template <class ValueT, class BaseT>
    KdTreeCell<ValueT, BaseT>::KdTreeCell()
    {
        memset(&_children, NULL, sizeof(_children));
    }
 
    template <class ValueT, class BaseT>
    KdTreeCell<ValueT, BaseT>::KdTreeCell(const KdTreeCell<ValueT, BaseT>& cell) :
        BaseT(cell), _split(cell._split), _box(cell._box)
    {
        if (cell._children[0])
        {
            _children[0] = new KdTreeCell<ValueT, BaseT>(*cell._children[0]);
        }
        else
        {
            _children[0] = NULL;
        }
        if (cell._children[1])
        {
            _children[1] = new KdTreeCell<ValueT, BaseT>(*cell._children[1]);
        }
        else
        {
            _children[1] = NULL;
        }
    }
 
    template <class ValueT, class BaseT>
    KdTreeCell<ValueT, BaseT>::~KdTreeCell()
    {
        for (unsigned int i = 0; i < 2; ++i)
            if (_children[i])
            {
                delete _children[i];
            }
    }
 
    template <class ValueT, class BaseT>
    const KdTreeCell<ValueT, BaseT>*
    KdTreeCell<ValueT, BaseT>::operator[](unsigned int index) const
    {
        return _children[index];
    }
 
    template <class ValueT, class BaseT>
    KdTreeCell<ValueT, BaseT>*
    KdTreeCell<ValueT, BaseT>::operator[](unsigned int index)
    {
        return _children[index];
    }
 
    template <class ValueT, class BaseT>
    void
    KdTreeCell<ValueT, BaseT>::Child(unsigned int index, ThisType* child)
    {
        _children[index] = child;
    }
 
    template <class ValueT, class BaseT>
    const typename KdTreeCell<ValueT, BaseT>::SplitterType&
    KdTreeCell<ValueT, BaseT>::Split() const
    {
        return _split;
    }
 
    template <class ValueT, class BaseT>
    const typename KdTreeCell<ValueT, BaseT>::BoxType&
    KdTreeCell<ValueT, BaseT>::Box() const
    {
        return _box;
    }
 
    template <class Strategies>
    void
    KdTree<Strategies>::Build()
    {
        Clear();
 
        CellType* root = new CellType;
        CellRange range;
        RootRange(&range);
        root->_box.Infinite();
        Root(root);
        Init(range, NULL, root);
 
        typedef std::pair<CellType*, CellRange> Pair;
        std::deque<Pair> stack;
        stack.push_back(Pair(root, range));
        while (stack.size())
        {
            Pair p = stack.back();
            stack.pop_back();
            if (ShouldSubdivide(p.second))
            {
                Subdivide(p.first, p.second);
                if (IsLeaf(*p.first)) // couldn't subdivide?
                {
                    continue;
                }
                for (unsigned int i = 0; i < CellType::NChildren; ++i)
                {
                    Range(*p.first, p.second, i, &range);
                    //Init(range, p.first, (*p.first)[i]);
                    stack.push_back(Pair((*p.first)[i], range));
                }
            }
        }
    }
 
    template <class Strategies>
    void
    KdTree<Strategies>::Insert(DereferencedType s)
    {
        CellType* cell = Root();
        CellRange range;
        RootRange(&range);
        if (!cell)
        {
            cell = new CellType;
            Root(cell);
            Init(range, NULL, cell);
        }
        Insert(cell, range, s);
    }
 
    template <class Strategies>
    template <class ContainerT>
    void
    KdTree<Strategies>::NearestNeighbors(const PointType& p,
                                         unsigned int k,
                                         ContainerT* neighbors,
                                         ScalarType* dist) const
    {
        *dist = std::numeric_limits<ScalarType>::infinity();
        neighbors->resize(0);
        if (!Root())
        {
            return;
        }
        AABox<PointType> rootBox;
        //rootBox.Infinite();
        CellRange range;
        RootRange(&range);
        neighbors->reserve(k + 1);
        NearestNeighbors(*Root(), range, rootBox, p, k, neighbors, dist);
        if (*dist < std::numeric_limits<ScalarType>::infinity())
        {
            *dist = std::sqrt(*dist);
        }
    }
 
    template <class Strategies>
    template <template <class> class ContainerT>
    void
    KdTree<Strategies>::NearestNeighborsL(const PointType& p,
                                          unsigned int k,
                                          LimitedHeap<NN, ContainerT>* neighbors,
                                          ScalarType* dist) const
    {
        *dist = std::numeric_limits<ScalarType>::infinity();
        neighbors->resize(0);
        if (!Root())
        {
            return;
        }
        neighbors->Limit(k);
        CellRange range;
        RootRange(&range);
        NearestNeighborsL(*Root(), range, p, k, neighbors, dist);
        //neighbors->sort_heap();
        //std::sort_heap(neighbors->begin(), neighbors->end());
        //if(*dist < std::numeric_limits< ScalarType >::infinity())
        //  *dist = std::sqrt(*dist);
    }
 
    template <class Strategies>
    template <class ContainerT>
    void
    KdTree<Strategies>::ApproximateNearestNeighbors(const PointType& p,
                                                    unsigned int k,
                                                    ScalarType eps,
                                                    ContainerT* neighbors,
                                                    ScalarType* sqrDist) const
    {
        *sqrDist = std::numeric_limits<ScalarType>::infinity();
        neighbors->resize(0);
        if (!Root())
        {
            return;
        }
        AABox<PointType> rootBox;
        //rootBox.Infinite();
        CellRange range;
        RootRange(&range);
        neighbors->reserve(k + 1);
        eps = 1 - eps;
        ApproximateNearestNeighbors(*Root(), range, rootBox, p, k, eps * eps, neighbors, sqrDist);
        //if(*dist < std::numeric_limits< ScalarType >::infinity())
        //  *dist = std::sqrt(*dist);
    }
 
    template <class Strategies>
    template <class ContainerT>
    void
    KdTree<Strategies>::PointsInSphere(const PointType& p,
                                       ScalarType sqrRadius,
                                       ContainerT* points) const
    {
        points->resize(0);
        if (!Root())
        {
            return;
        }
        CellRange range;
        RootRange(&range);
        PointsInSphere(*Root(), range, p, sqrRadius, points);
    }
 
    template <class Strategies>
    bool
    KdTree<Strategies>::Contains(const PointType& p, DereferencedType* d) const
    {
        if (!Root())
        {
            return false;
        }
        CellRange range;
        RootRange(&range);
        return Contains(*Root(), range, p, d);
    }
 
    template <class Strategies>
    bool
    KdTree<Strategies>::ShouldSubdivide(CellRange range) const
    {
        return range.second - range.first > 10;
    }
 
    template <class Strategies>
    void
    KdTree<Strategies>::Subdivide(CellType* cell, CellRange range)
    {
        PointType pmin = cell->_box.Min(), pmax = cell->_box.Max();
        /*pmax = pmin = at(Dereference(range.first));
        for(HandleType i = range.first + 1; i < range.second; ++i)
        {
            const PointType p = at(Dereference(i));
            for(unsigned int j = 0; j < Dim; ++j)
            {
                if(pmin[j] > p[j])
                    pmin[j] = p[j];
                else if(pmax[j] < p[j])
                    pmax[j] = p[j];
            }
        }*/
 
        PointType pdiff = pmax - pmin;
        unsigned int axis = 0;
        ScalarType length = pdiff[0];
        for (unsigned int j = 1; j < Dim; ++j)
        {
            if (pdiff[j] > length)
            {
                axis = j;
                length = pdiff[j];
            }
        }
 
        ScalarType split = (pmax[axis] + pmin[axis]) / 2;
 
        cell->_split.Set(axis, split);
 
        CellType *left, *right;
        left = new CellType;
        right = new CellType;
        Split(cell->_split, range, left, right);
        if (left->Size() == cell->Size() || right->Size() == cell->Size())
        {
            delete left;
            delete right;
            return;
        }
        //cell->_box.Split(cell->_split.Axis(), cell->_split.Value(),
        //  &left->_box, &right->_box);
        cell->Child(0, left);
        cell->Child(1, right);
        CellRange childRange;
        Range(*cell, range, 0, &childRange);
        Init(childRange, cell, left);
        Range(*cell, range, 1, &childRange);
        Init(childRange, cell, right);
        // refine split value
        cell->_split.Value(
            (right->Box().Min()[cell->_split.Axis()] + left->Box().Max()[cell->_split.Axis()]) / 2);
    }
 
    template <class Strategies>
    void
    KdTree<Strategies>::Insert(CellType* cell, CellRange range, DereferencedType s)
    {
        cell->_box += at(s);
        if (IsLeaf(*cell))
        {
            BaseType::Insert(range, s);
            if (ShouldSubdivide(range))
            {
                Subdivide(cell, range);
            }
            return;
        }
        bool left;
        BaseType::SplitAndInsert(cell->Split(), range, s, (*cell)[0], (*cell)[1], &left);
        if (left)
        {
            CellRange lr;
            Range(*cell, range, 0, &lr);
            Insert((*cell)[0], lr, s);
        }
        else
        {
            CellRange rr;
            Range(*cell, range, 1, &rr);
            Insert((*cell)[1], rr, s);
        }
    }
 
    template <class Strategies>
    template <class ContainerT>
    void
    KdTree<Strategies>::NearestNeighbors(const CellType& cell,
                                         CellRange range,
                                         const AABox<PointType>& box,
                                         const PointType& p,
                                         unsigned int k,
                                         ContainerT* neighbors,
                                         ScalarType* dist2) const
    {
        if (IsLeaf(cell))
        {
            HandleType i = range.first, iend = range.second;
            if (neighbors->size() + cell.Size() <= k)
            {
                for (; i != iend; ++i)
                {
                    DereferencedType deref = Dereference(i);
                    PointType diff = at(deref);
                    diff -= p;
                    ScalarType d2 = diff * diff;
                    neighbors->push_back(NN(deref, d2));
                }
                std::sort(neighbors->begin(), neighbors->end());
                *dist2 = neighbors->back().sqrDist;
                return;
            }
            // let's find nearest neighbors from points within the cell
            for (; i != iend; ++i)
            {
                DereferencedType deref = Dereference(i);
                PointType diff = at(deref);
                diff -= p;
                ScalarType d2 = diff * diff;
                if (neighbors->size() < k || d2 <= *dist2)
                {
                    SortedNearestNeighborInsert(NN(deref, d2), k, neighbors);
                    *dist2 = neighbors->back().sqrDist;
                }
            }
            return;
        }
        // descent into subtree containing p
        unsigned int nearerChild, fartherChild;
        if (cell.Split()(p) <= 0)
        {
            nearerChild = 0;
            fartherChild = 1;
        }
        else
        {
            nearerChild = 1;
            fartherChild = 0;
        }
        /*AABox< PointType > subBoxes[2];
        box.Split(cell.Split().Axis(), cell.Split().Value(),
            &subBoxes[0], &subBoxes[1]);*/
 
        CellRange childRange;
        if (neighbors->size() < k || Intersect(cell[nearerChild]->Box(), p, *dist2))
        {
            Range(cell, range, nearerChild, &childRange);
            NearestNeighbors(*cell[nearerChild],
                             childRange,
                             box /*subBoxes[nearerChild]*/,
                             p,
                             k,
                             neighbors,
                             dist2);
        }
 
        // check if closer neighbor could be in other child
        if (neighbors->size() < k ||
            Intersect(cell[fartherChild]->Box() /*subBoxes[fartherChild]*/, p, *dist2))
        {
            Range(cell, range, fartherChild, &childRange);
            NearestNeighbors(*cell[fartherChild],
                             childRange,
                             box /*subBoxes[fartherChild]*/,
                             p,
                             k,
                             neighbors,
                             dist2);
        }
    }
 
    template <class Strategies>
    template <template <class> class ContainerT>
    void
    KdTree<Strategies>::NearestNeighborsL(const CellType& cell,
                                          const CellRange& range,
                                          const PointType& p,
                                          unsigned int k,
                                          LimitedHeap<NN, ContainerT>* neighbors,
                                          ScalarType* dist2) const
    {
        if (IsLeaf(cell))
        {
            HandleType i = range.first, iend = range.second;
            DereferencedType deref;
            if (neighbors->size() + cell.Size() <= k)
            {
                for (; i != iend; ++i)
                {
                    deref = Dereference(i);
                    neighbors->push_back(NN(deref, p.SqrDistance(at(deref).Point())));
                }
                neighbors->MakeHeap();
                *dist2 = neighbors->front().sqrDist;
                return;
            }
            // let's find nearest neighbors from points within the cell
            for (; i != iend; ++i)
            {
                deref = Dereference(i);
                neighbors->PushHeap(NN(deref, p.SqrDistance(at(deref).Point())));
                *dist2 = neighbors->front().sqrDist;
            }
            return;
        }
        // descent into subtree containing p
        unsigned int nearerChild, fartherChild;
        const CellType *nChild, *fChild;
        ScalarType splitDist = cell._split(p);
        if (splitDist <= 0)
        {
            nearerChild = 0;
            fartherChild = 1;
            nChild = cell[0];
            fChild = cell[1];
        }
        else
        {
            nearerChild = 1;
            fartherChild = 0;
            nChild = cell[1];
            fChild = cell[0];
        }
 
        CellRange childRange;
        //if(neighbors->size() < k || Intersect(p, *dist2, nChild->_box.Min(),
        //  nChild->_box.Max()))
        //{
        Range(cell, range, nearerChild, &childRange);
        NearestNeighborsL(*nChild, childRange, p, k, neighbors, dist2);
        //}
 
        // check if closer neighbor could be in other child
        if (neighbors->size() < k || (splitDist * splitDist < *dist2 &&
                                      Intersect(p, *dist2, fChild->_box.Min(), fChild->_box.Max())))
        {
            Range(cell, range, fartherChild, &childRange);
            NearestNeighborsL(*fChild, childRange, p, k, neighbors, dist2);
        }
    }
 
    template <class Strategies>
    template <class ContainerT>
    void
    KdTree<Strategies>::ApproximateNearestNeighbors(const CellType& cell,
                                                    CellRange range,
                                                    const AABox<PointType>& box,
                                                    const PointType& p,
                                                    unsigned int k,
                                                    ScalarType eps,
                                                    ContainerT* neighbors,
                                                    ScalarType* sqrDist) const
    {
        if (IsLeaf(cell))
        {
            if (neighbors->size() + cell.Size() <= k)
            {
                for (HandleType i = range.first; i != range.second; ++i)
                {
                    DereferencedType deref = Dereference(i);
                    PointType diff = at(deref);
                    diff -= p;
                    ScalarType d2 = diff * diff;
                    neighbors->push_back(NN(deref, d2));
                }
                std::sort(neighbors->begin(), neighbors->end());
                *sqrDist = neighbors->back().sqrDist;
                return;
            }
            // let's find nearest neighbors from points within the cell
            for (HandleType i = range.first; i != range.second; ++i)
            {
                DereferencedType deref = Dereference(i);
                PointType diff = at(deref);
                diff -= p;
                ScalarType d2 = diff * diff;
                if (neighbors->size() < k || d2 <= *sqrDist)
                {
                    SortedNearestNeighborInsert(NN(deref, d2), k, neighbors);
                    *sqrDist = neighbors->back().sqrDist;
                }
            }
            return;
        }
        // descent into subtree containing p
        unsigned int nearerChild, fartherChild;
        if (cell.Split()(p) <= 0)
        {
            nearerChild = 0;
            fartherChild = 1;
        }
        else
        {
            nearerChild = 1;
            fartherChild = 0;
        }
        /*AABox< PointType > subBoxes[2];
        box.Split(cell.Split().Axis(), cell.Split().Value(),
            &subBoxes[0], &subBoxes[1]);*/
 
        CellRange childRange;
        if (neighbors->size() < k || Intersect(cell[nearerChild]->Box(), p, eps * (*sqrDist)))
        {
            Range(cell, range, nearerChild, &childRange);
            ApproximateNearestNeighbors(*cell[nearerChild],
                                        childRange,
                                        box /*subBoxes[nearerChild]*/,
                                        p,
                                        k,
                                        eps,
                                        neighbors,
                                        sqrDist);
        }
 
        // check if closer neighbor could be in other child
        if (neighbors->size() < k ||
            Intersect(cell[fartherChild]->Box() /*subBoxes[fartherChild]*/, p, eps * (*sqrDist)))
        {
            Range(cell, range, fartherChild, &childRange);
            ApproximateNearestNeighbors(*cell[fartherChild],
                                        childRange,
                                        box /*subBoxes[fartherChild]*/,
                                        p,
                                        k,
                                        eps,
                                        neighbors,
                                        sqrDist);
        }
    }
 
    template <class Strategies>
    template <class ContainerT>
    void
    KdTree<Strategies>::PointsInSphere(const CellType& cell,
                                       const CellRange& range,
                                       const PointType& p,
                                       ScalarType sqrRadius,
                                       ContainerT* points) const
    {
        if (IsLeaf(cell))
        {
            for (HandleType i = range.first; i < range.second; ++i)
            {
                DereferencedType deref = Dereference(i);
                ScalarType sqrDist = p.SqrDistance(at(deref));
                if (sqrDist <= sqrRadius)
                {
                    points->push_back(NN(deref, sqrDist));
                }
            }
        }
        else
        {
            CellRange childRange;
            if (Intersect(p, sqrRadius, cell[0]->_box.Min(), cell[0]->_box.Max()))
            {
                Range(cell, range, 0, &childRange);
                PointsInSphere(*(cell[0]), childRange, p, sqrRadius, points);
            }
            if (Intersect(p, sqrRadius, cell[1]->_box.Min(), cell[1]->_box.Max()))
            {
                Range(cell, range, 1, &childRange);
                PointsInSphere(*(cell[1]), childRange, p, sqrRadius, points);
            }
        }
    }
 
    template <class Strategies>
    bool
    KdTree<Strategies>::Contains(const CellType& cell,
                                 const CellRange range,
                                 const PointType& p,
                                 DereferencedType* d) const
    {
        if (IsLeaf(cell))
        {
            for (HandleType i = range.first; i < range.second; ++i)
            {
                DereferencedType deref = Dereference(i);
                if (p == at(deref))
                {
                    *d = deref;
                    return true;
                }
            }
            return false;
        }
        // descent into subtree containing p
        unsigned int nearerChild = cell.Split()(p) <= 0 ? 0 : 1;
        CellRange childRange;
        Range(cell, range, nearerChild, &childRange);
        return Contains(*cell[nearerChild], childRange, p, d);
    }
 
    template <class Strategies>
    void
    KdTree<Strategies>::Init(const CellRange& range, const CellType*, CellType* cell)
    {
        PointType pmin, pmax;
        pmax = pmin = at(Dereference(range.first));
        for (HandleType i = range.first + 1; i != range.second; ++i)
        {
            const PointType p = at(Dereference(i));
            for (unsigned int j = 0; j < Dim; ++j)
            {
                if (pmin[j] > p[j])
                {
                    pmin[j] = p[j];
                }
                else if (pmax[j] < p[j])
                {
                    pmax[j] = p[j];
                }
            }
        }
        cell->_box.Min() = pmin;
        cell->_box.Max() = pmax;
    }
}; // namespace GfxTL