database.h

Go to the documentation of this file.

#ifndef VOCABULARY_TREE_DATABASE_H
#define VOCABULARY_TREE_DATABASE_H

#include <vocabulary_tree/vocabulary_tree.h>
#include <map>

namespace vt {

typedef uint32_t DocId;

struct Match
{
  DocId id;
  float score;

  Match() {}
  Match(DocId _id, float _score) : id(_id), score(_score) {}

  bool operator<(const Match& other) const
  {
    return score < other.score;
  }
};

// Remove these, just make docs more confusing
typedef std::vector<Word> Document;
typedef std::vector<Match> Matches;

class Database
{
public:
  Database(uint32_t num_words = 0);

  DocId insert(const std::vector<Word>& document);

  void find(const std::vector<Word>& document, size_t N, std::vector<Match>& matches) const;

  DocId findAndInsert(const std::vector<Word>& document, size_t N, std::vector<Match>& matches);

  void computeTfIdfWeights(float default_weight = 1.0f);

  void saveWeights(const std::string& file) const;
  void loadWeights(const std::string& file);

  // Save weights and documents
  //void save(const std::string& file) const;
  //void load(const std::string& file);

private:
  struct WordFrequency
  {
    DocId id;
    uint32_t count;

    WordFrequency(DocId _id, uint32_t _count) : id(_id), count(_count) {}
  };
  
  // Stored in increasing order by DocId
  typedef std::vector<WordFrequency> InvertedFile;

  // typedef std::vector< std::pair<Word, float> > DocumentVector;
  typedef std::map<Word, float> DocumentVector;

  std::vector<InvertedFile> word_files_;
  std::vector<float> word_weights_;
  std::vector<DocumentVector> database_vectors_; // Precomputed for inserted documents

  void computeVector(const std::vector<Word>& document, DocumentVector& v) const;
  
  static void normalize(DocumentVector& v);
  static float sparseDistance(const DocumentVector& v1, const DocumentVector& v2);
};

} //namespace vt

#endif