#include <cstdlib>
#include <list>
#include <vector>
#include <algorithm>
#include <set>
#include "utility.h"
namespace meow{
template<class Keys, class Key, class Value>
inline KD_Tree<Keys,Key,Value>::Node::Node(Keys _key,
Value _value,
ssize_t _l_child,
ssize_t _r_child):
key(_key), value(_value), lChild(_l_child), rChild(_r_child){ }
//
template<class Keys, class Key, class Value>
inline KD_Tree<Keys, Key, Value>::Sorter::Sorter(KD_Tree::Nodes const& _nodes,
size_t _cmp):
nodes(_nodes), cmp(_cmp){ }
template<class Keys, class Key, class Value>
inline bool
KD_Tree<Keys, Key, Value>::Sorter::operator()(size_t const& a,
size_t const& b) const{
if(nodes[a].key[cmp] != nodes[b].key[cmp]){
return (nodes[a].key[cmp] < nodes[b].key[cmp]);
}
return (nodes[a].value < nodes[b].value);
}
//
template<class Keys, class Key, class Value>
inline KD_Tree<Keys, Key, Value>::Answer::Answer(Node const& _node,
Key _dist2):
node(_node), dist2(_dist2){ }
template<class Keys, class Key, class Value>
inline bool
KD_Tree<Keys, Key, Value>::Answer::operator<(Answer const& b) const{
if(dist2 != b.dist2) return (dist2 < b.dist2);
return (node.value < b.node.value);
}
//
template<class Keys, class Key, class Value>
inline Key KD_Tree<Keys, Key, Value>::distance2(Keys const& k1,
Keys const& k2) const{
Key ret(0);
for(size_t i = 0; i < dimension; i++)
ret += squ(k1[i] - k2[i]);
return ret;
}
template<class Keys, class Key, class Value>
inline size_t KD_Tree<Keys, Key, Value>::query(Keys const& key,
size_t k,
size_t index,
int depth,
std::vector<Key>& dist2_v,
Key dist2_s,
KD_Tree::AnswerList* ret) const{
if(index == NIL){
return 0;
}
size_t cmp = depth % dimension;
ssize_t right_side, opposite, size;
ssize_t sz, other;
if(key[cmp] <= nodes[index].key[cmp]){
right_side = nodes[index].lChild;
opposite = nodes[index].rChild;
}else{
right_side = nodes[index].rChild;
opposite = nodes[index].lChild;
}
size = query(key, k, right_side, depth + 1, dist2_v, dist2_s, ret);
Answer my_ans(nodes[index], distance2(nodes[index].key, key));
if(size < k || my_ans < *(ret->rbegin())){
KD_Tree::AnswerListIterator it = ret->begin();
while(it != ret->end() && !(my_ans < *it)) it++;
ret->insert(it, my_ans);
size++;
}
Key dist2_old = dist2_v[cmp];
dist2_v[cmp] = squ(nodes[index].key[cmp] - key[cmp]);
dist2_s += dist2_v[cmp] - dist2_old;
if(size < k || (*(ret->rbegin())).dist2 >= dist2_s){
KD_Tree::AnswerList ret2;
size += query(key, k, opposite, depth + 1, dist2_v, dist2_s, &ret2);
KD_Tree::AnswerListIterator it1, it2;
for(it1 = ret->begin(), it2 = ret2.begin(); it2 != ret2.end(); it2++){
while(it1 != ret->end() && *it1 < *it2) it1++;
it1 = ++(ret->insert(it1, *it2));
}
}
if(size > k){
for(int i = size - k; i--; ){
ret->pop_back();
}
size = k;
}
dist2_v[cmp] = dist2_old;
return size;
}
template<class Keys, class Key, class Value>
inline ssize_t KD_Tree<Keys, Key, Value>::build(ssize_t beg,
ssize_t end,
std::vector<size_t>* orders,
int depth){
if(beg > end){
return NIL;
}
ssize_t mid = (beg + end) / 2;
size_t cmp = depth % 2;
std::set<size_t> right;
for(ssize_t i = mid + 1; i <= end; i++){
right.insert(orders[cmp][i]);
}
for(int i = 0; i < dimension; i++){
if(i == cmp) continue;
size_t aa = beg, bb = mid + 1;
for(int j = beg; j <= end; j++){
if(orders[i][j] == orders[cmp][mid]){
orders[dimension][mid] = orders[i][j];
}else if(right.find(orders[i][j]) != right.end()){
orders[dimension][bb++] = orders[i][j];
}else{
orders[dimension][aa++] = orders[i][j];
}
}
for(int j = beg; j <= end; j++){
orders[i][j] = orders[dimension][j];
}
}
nodes[orders[cmp][mid]].lChild = build(beg, mid - 1, orders, depth + 1);
nodes[orders[cmp][mid]].rChild = build(mid + 1, end, orders, depth + 1);
return orders[cmp][mid];
}
template<class Keys, class Key, class Value>
inline KD_Tree<Keys, Key, Value>::KD_Tree():
NIL(-1), root(NIL), needRebuild(false), dimension(1){ }
template<class Keys, class Key, class Value>
inline KD_Tree<Keys, Key, Value>::KD_Tree(size_t _dimension):
NIL(-1), root(NIL), needRebuild(false), dimension(_dimension){ }
template<class Keys, class Key, class Value>
inline KD_Tree<Keys, Key, Value>::~KD_Tree(){ }
template<class Keys, class Key, class Value>
inline void KD_Tree<Keys, Key, Value>::insert(Keys const& key, Value value){
nodes.push_back(Node(key, value, NIL, NIL));
needRebuild = true;
}
template<class Keys, class Key, class Value>
inline void KD_Tree<Keys, Key, Value>::build(){
if(needRebuild){
std::vector<size_t> *orders = new std::vector<size_t>[dimension + 1];
for(int j = 0; j < dimension + 1; j++){
orders[j].resize(nodes.size());
}
for(int j = 0; j < dimension; j++){
for(size_t i = 0, I = nodes.size(); i < I; i++){
orders[j][i] = i;
}
std::sort(orders[j].begin(), orders[j].end(), Sorter(nodes, j));
}
root = build(0, (ssize_t)nodes.size() - 1, orders, 0);
needRebuild = false;
delete [] orders;
}
}
template<class Keys, class Key, class Value>
inline Value KD_Tree<Keys, Key, Value>::query(Keys const& point, int k) const{
((KD_Tree*)this)->build();
KD_Tree::AnswerList ret;
std::vector<Key> tmp(dimension, Key(0));
query(point, k, root, 0, tmp, Key(0), &ret);
return (*(ret.rbegin())).node.value;
}
template<class Keys, class Key, class Value>
inline typename KD_Tree<Keys, Key, Value>::Values
KD_Tree<Keys, Key, Value>::rangeQuery(Keys const& point, int k) const{
((KD_Tree*)this)->build();
KD_Tree::AnswerList ret;
std::vector<Key> tmp(dimension, Key(0));
query(point, k, root, 0, tmp, Key(0), &ret);
KD_Tree::Values ret_val(ret.size());
int i = 0;
for(KD_Tree::AnswerListIterator it = ret.begin(); it != ret.end(); it++){
ret_val[i++] = (*it).node.value;
}
return ret_val;
}
template<class Keys, class Key, class Value>
inline void KD_Tree<Keys, Key, Value>::clear(){
root = NIL;
nodes.clear();
needRebuild = false;
}
template<class Keys, class Key, class Value>
inline void KD_Tree<Keys, Key, Value>::reset(size_t _dimension){
clear();
dimension = _dimension;
}
}
