forked from ZJULearning/nsg
-
Notifications
You must be signed in to change notification settings - Fork 0
/
index_nsg.h
85 lines (70 loc) · 2.49 KB
/
index_nsg.h
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
#ifndef EFANNA2E_INDEX_NSG_H
#define EFANNA2E_INDEX_NSG_H
#include "util.h"
#include "parameters.h"
#include "neighbor.h"
#include "index.h"
#include <cassert>
#include <unordered_map>
#include <string>
#include <sstream>
#include <boost/dynamic_bitset.hpp>
#include <stack>
namespace efanna2e {
class IndexNSG : public Index {
public:
explicit IndexNSG(const size_t dimension, const size_t n, Metric m, Index *initializer);
virtual ~IndexNSG();
virtual void Save(const char *filename)override;
virtual void Load(const char *filename)override;
virtual void Build(size_t n, const float *data, const Parameters ¶meters) override;
virtual void Search(
const float *query,
const float *x,
size_t k,
const Parameters ¶meters,
unsigned *indices) override;
void SearchWithOptGraph(
const float *query,
size_t K,
const Parameters ¶meters,
unsigned *indices);
void OptimizeGraph(float* data);
protected:
typedef std::vector<std::vector<unsigned > > CompactGraph;
typedef std::vector<SimpleNeighbors > LockGraph;
typedef std::vector<nhood> KNNGraph;
CompactGraph final_graph_;
Index *initializer_;
void init_graph(const Parameters ¶meters);
void get_neighbors(
const float *query,
const Parameters ¶meter,
std::vector<Neighbor> &retset,
std::vector<Neighbor> &fullset);
void get_neighbors(
const float *query,
const Parameters ¶meter,
boost::dynamic_bitset<>& flags,
std::vector<Neighbor> &retset,
std::vector<Neighbor> &fullset);
//void add_cnn(unsigned des, Neighbor p, unsigned range, LockGraph& cut_graph_);
void InterInsert(unsigned n, unsigned range, std::vector<std::mutex>& locks, SimpleNeighbor* cut_graph_);
void sync_prune(unsigned q, std::vector<Neighbor>& pool, const Parameters ¶meter, boost::dynamic_bitset<>& flags, SimpleNeighbor* cut_graph_);
void Link(const Parameters ¶meters, SimpleNeighbor* cut_graph_);
void Load_nn_graph(const char *filename);
void tree_grow(const Parameters ¶meter);
void DFS(boost::dynamic_bitset<> &flag, unsigned root, unsigned &cnt);
void findroot(boost::dynamic_bitset<> &flag, unsigned &root, const Parameters ¶meter);
private:
unsigned width;
unsigned ep_;
std::vector<std::mutex> locks;
char* opt_graph_;
size_t node_size;
size_t data_len;
size_t neighbor_len;
KNNGraph nnd_graph;
};
}
#endif //EFANNA2E_INDEX_NSG_H