knn.h

Go to the documentation of this file.

#ifndef META_KNN_H_
#define META_KNN_H_

#include <unordered_set>
#include "index/inverted_index.h"
#include "index/forward_index.h"
#include "index/ranker/ranker.h"
#include "classify/classifier_factory.h"
#include "classify/classifier/classifier.h"

namespace meta
{
namespace classify
{

class knn : public classifier
{
  public:
    const static std::string id;

    knn(std::shared_ptr<index::inverted_index> idx,
        std::shared_ptr<index::forward_index> f_idx, uint16_t k,
        std::unique_ptr<index::ranker> ranker, bool weighted = false);

    void train(const std::vector<doc_id>& docs) override;

    class_label classify(doc_id d_id) override;

    void reset() override;

  private:
    class_label select_best_label(
        const std::vector<std::pair<doc_id, double>>& scored,
        const std::vector<std::pair<class_label, uint16_t>>& sorted) const;

    std::shared_ptr<index::inverted_index> inv_idx_;

    uint16_t k_;

    std::unique_ptr<index::ranker> ranker_;

    std::unordered_set<doc_id> legal_docs_;

    const bool weighted_;

  public:
    class knn_exception : public std::runtime_error
    {
      public:
        using std::runtime_error::runtime_error;
    };
};

template <>
std::unique_ptr<classifier>
    make_multi_index_classifier<knn>(const cpptoml::table&,
                                     std::shared_ptr<index::forward_index>,
                                     std::shared_ptr<index::inverted_index>);
}
}
#endif