NeighborQueue.hpp

Go to the documentation of this file.

#ifndef KNNCOLLE_NEIGHBOR_QUEUE_HPP
#define KNNCOLLE_NEIGHBOR_QUEUE_HPP
 
#include "utils.hpp"
 
#include <queue>
#include <vector>
#include <algorithm>
#include <cassert>
 
#include "sanisizer/sanisizer.hpp"
 
namespace knncolle {
 
template<typename Index_, typename Distance_>
class NeighborQueue {
public:
    NeighborQueue() = default;
 
public:    
    void reset(Index_ k) {
        // We don't allow k == 0 as otherwise we'd be in a position where is_full() == true but limit() can't be called.
        // If the caller doesn't want any neighbors, they're better of just aborting the search altogether.
        assert(k > 0);
        my_neighbors = sanisizer::cast<I<decltype(my_neighbors)> >(k);
        my_full = false;
 
        // Popping any existing elements out, just in case. This shouldn't
        // usually be necessary if report() was called as the queue should
        // already be completely exhausted, but sometimes report() is a no-op
        // or there might have been an intervening exception, etc.
        while (!my_nearest.empty()) {
            my_nearest.pop();
        }
    }
 
public:
    bool is_full() const {
        return my_full;
    }
 
    Distance_ limit() const {
        return my_nearest.top().first;
   }
 
    auto size() const {
        return my_nearest.size();
    }
 
public:
    void add(Index_ i, Distance_ d) {
        my_nearest.emplace(d, i);
        if (my_full) {
            my_nearest.pop();
        } else if (size() == my_neighbors) {
            my_full = true;
        }
    }
 
private:
    template<bool has_indices_, bool has_distances_>
    void report_internal(std::vector<Index_>* output_indices, std::vector<Distance_>* output_distances, const Index_ self) {
        // We expect that nearest is non-empty, as a search should at least find 'self' (or duplicates thereof).
        assert(!my_nearest.empty());
        const Index_ num_expected = size() - 1;
 
        if constexpr(has_indices_) {
            output_indices->clear();
            output_indices->reserve(num_expected);
        }
        if constexpr(has_distances_) {
            output_distances->clear();
            output_distances->reserve(num_expected);
        }
 
        bool found_self = false;
        while (!my_nearest.empty()) {
            const auto& top = my_nearest.top();
            if (!found_self && top.second == self) {
                found_self = true;
            } else {
                if constexpr(has_indices_) {
                    output_indices->push_back(top.second);
                }
                if constexpr(has_distances_) {
                    output_distances->push_back(top.first);
                }
            }
            my_nearest.pop();
        }
 
        // We use push_back + reverse to give us sorting in increasing order;
        // this is nicer than push_front() for std::vectors.
        if constexpr(has_indices_) {
            std::reverse(output_indices->begin(), output_indices->end());
        }
        if constexpr(has_distances_) {
            std::reverse(output_distances->begin(), output_distances->end());
        }
 
        // Removing the most distance element if we couldn't find ourselves,
        // e.g., because there are too many duplicates.
        if (!found_self) {
            if constexpr(has_indices_) {
                output_indices->pop_back();
            }
            if constexpr(has_distances_) {
                output_distances->pop_back();
            }
        }
    } 
 
public:
    void report(std::vector<Index_>* output_indices, std::vector<Distance_>* output_distances, Index_ self) {
        if (output_indices && output_distances) {
            report_internal<true, true>(output_indices, output_distances, self);
        } else if (output_indices) {
            report_internal<true, false>(output_indices, NULL, self);
        } else if (output_distances) {
            report_internal<false, true>(NULL, output_distances, self);
        }
    }
 
private:
    template<bool has_indices_, bool has_distances_>
    void report_internal(std::vector<Index_>* output_indices, std::vector<Distance_>* output_distances) {
        auto position = my_nearest.size();
 
        if constexpr(has_indices_) {
            sanisizer::resize(*output_indices, position);
        }
        if constexpr(has_distances_) {
            sanisizer::resize(*output_distances, position);
        }
 
        while (!my_nearest.empty()) {
            const auto& top = my_nearest.top();
            --position;
            if constexpr(has_indices_) {
                (*output_indices)[position] = top.second;
            }
            if constexpr(has_distances_) {
                (*output_distances)[position] = top.first;
            }
            my_nearest.pop();
        }
    } 
 
public:
    void report(std::vector<Index_>* output_indices, std::vector<Distance_>* output_distances) {
        if (output_indices && output_distances) {
            report_internal<true, true>(output_indices, output_distances);
        } else if (output_indices) {
            report_internal<true, false>(output_indices, NULL);
        } else if (output_distances) {
            report_internal<false, true>(NULL, output_distances);
        }
    }
 
private:
    bool my_full = false;
    std::priority_queue<std::pair<Distance_, Index_> > my_nearest;
    I<decltype(my_nearest.size())> my_neighbors = 1;
};
 
}
 
#endif