330 lines
10 KiB
C++
330 lines
10 KiB
C++
#include <Eigen/Dense>
|
|
#include <algorithm>
|
|
#include <array>
|
|
#include <bitset>
|
|
#include <deque>
|
|
#include <execution>
|
|
#include <forward_list>
|
|
#include <functional>
|
|
#include <iostream>
|
|
#include <list>
|
|
#include <map>
|
|
#include <memory>
|
|
#include <mutex>
|
|
#include <queue>
|
|
#include <random>
|
|
#include <shared_mutex>
|
|
#include <tuple>
|
|
#include <unordered_set>
|
|
|
|
template <unsigned int dim, typename number, typename T> class Orthtree {
|
|
public:
|
|
typedef Eigen::Vector<number, dim> Point;
|
|
typedef std::function<bool(std::unique_ptr<T> &, const Point &)>
|
|
data_access_f;
|
|
typedef std::function<bool(const std::unique_ptr<T> &, const Point &)>
|
|
cdata_access_f;
|
|
|
|
class BBox {
|
|
public:
|
|
enum Relation {
|
|
NO_OVERLAP,
|
|
OVERLAP,
|
|
CONTAINS,
|
|
CONTAINED,
|
|
};
|
|
Point min, max;
|
|
Point center() { return (min + max) / 2; }
|
|
Relation relation(const BBox &b) {
|
|
if ((b.min.array() > max.array()).any() ||
|
|
(b.max.array() < min.array()).any())
|
|
return NO_OVERLAP;
|
|
if ((b.min.array() > min.array()).all() &&
|
|
(b.max.array() < max.array()).all())
|
|
return CONTAINS;
|
|
if ((b.min.array() < min.array()).all() &&
|
|
(b.max.array() > max.array()).all())
|
|
return CONTAINED;
|
|
return OVERLAP;
|
|
};
|
|
BBox(const Point &min, const Point &max) : min(min), max(max) {
|
|
assert((min.array() < max.array()).all());
|
|
}
|
|
};
|
|
typedef std::bitset<dim> coord;
|
|
class Node;
|
|
struct NodeInfo {
|
|
const Node *node;
|
|
int depth;
|
|
number width;
|
|
coord crd;
|
|
};
|
|
typedef std::forward_list<NodeInfo> branch;
|
|
class Node {
|
|
public:
|
|
Point center;
|
|
std::unique_ptr<std::array<Node, (1 << dim)>> children;
|
|
std::unique_ptr<T> data;
|
|
|
|
mutable std::shared_mutex mtx;
|
|
|
|
void create_children(const Point &parent_center, number parent_width) {
|
|
children = std::make_unique<std::array<Node, (1 << dim)>>();
|
|
for (unsigned int i = 0; i < (1 << dim); ++i) {
|
|
Point offset = Point::Ones() * (parent_width / 4);
|
|
for (unsigned int j = 0; j < dim; ++j)
|
|
if (!coord(i)[j])
|
|
offset(j) *= -1;
|
|
(*children)[i].center = parent_center + offset;
|
|
}
|
|
};
|
|
|
|
coord get_child_coord(const Point &p) const {
|
|
coord child_coord;
|
|
for (unsigned int i = 0; i < dim; ++i)
|
|
child_coord[i] = center[i] < p[i];
|
|
return child_coord;
|
|
};
|
|
|
|
void depth(const Point &point, int depth, data_access_f f, number width) {
|
|
if (depth == 0) {
|
|
std::unique_lock guard(mtx);
|
|
f(data, center);
|
|
} else {
|
|
mtx.lock_shared();
|
|
if (!children) {
|
|
mtx.unlock_shared();
|
|
{
|
|
std::unique_lock uguard(mtx);
|
|
if (!children)
|
|
create_children(center, width);
|
|
}
|
|
mtx.lock_shared();
|
|
}
|
|
(*children)[get_child_coord(point).to_ulong()].depth(point, depth - 1,
|
|
f, width / 2);
|
|
|
|
mtx.unlock_shared();
|
|
}
|
|
};
|
|
|
|
BBox bb(number width) const {
|
|
Point offset = Point::Ones() * (width / 2);
|
|
return BBox(center - offset, center + offset);
|
|
};
|
|
|
|
void access_data(cdata_access_f f) const {
|
|
std::unique_lock guard(mtx);
|
|
if (data) {
|
|
f(data, center);
|
|
}
|
|
}
|
|
bool access_data_ret(cdata_access_f f, bool &ret) const {
|
|
std::unique_lock guard(mtx);
|
|
if (data) {
|
|
ret = f(data, center);
|
|
return true;
|
|
}
|
|
return false;
|
|
}
|
|
void all_data(cdata_access_f f) const {
|
|
std::unique_lock guard(mtx);
|
|
if (data) {
|
|
f(data, center);
|
|
}
|
|
if (children) {
|
|
std::for_each(std::execution::par_unseq, std::begin(*children),
|
|
std::end(*children),
|
|
[f](const auto &c) { c.all_data(f); });
|
|
}
|
|
}
|
|
bool is_leaf() const {
|
|
std::shared_lock guard(mtx);
|
|
return !bool(children);
|
|
}
|
|
void nodes_to_leaf(const Point &p, const number width, const int depth,
|
|
branch &l) const {
|
|
std::shared_lock guard(mtx);
|
|
if (!children)
|
|
return;
|
|
auto c = get_child_coord(p);
|
|
auto child = &(*children)[c.to_ullong()];
|
|
l.push_front({child, depth - 1, width / 2, c});
|
|
child->nodes_to_leaf(p, width / 2, depth - 1, l);
|
|
}
|
|
};
|
|
std::unique_ptr<Node> root;
|
|
number root_width;
|
|
|
|
mutable std::mutex mtx;
|
|
Orthtree(const Point &root_center, number root_width)
|
|
: root(new Node()), root_width(root_width) {
|
|
root->center = root_center;
|
|
};
|
|
|
|
void all_data(cdata_access_f f) const { root->all_data(f); };
|
|
|
|
void depth(const Point &point, int depth, data_access_f f) {
|
|
root->depth(point, depth, f, root_width);
|
|
}
|
|
|
|
// std::vector<NodeInfo> const std::forward_list<std::pair<Node *, coord>>&
|
|
// branch)
|
|
|
|
branch get_branch_to_leaf(const Point &p) const {
|
|
branch b({root.get, 0, root_width, coord{0}});
|
|
root->node_to_leaf(p, root_width, b);
|
|
return b;
|
|
}
|
|
|
|
// std::vector<branch> get_leaf_neighours(branch b) {
|
|
// auto node = b.front();
|
|
// b.pop_front();
|
|
// if (b.empty())
|
|
// return {};
|
|
// std::vector<branch> ret;
|
|
// ret.reserve(2<<dim);
|
|
// for (int i = 0; i < (1<<dim); i += 1) {
|
|
// auto& c = *b.front()->children;
|
|
// if (coord{i} != node->coord) {
|
|
// ret.push_back(b);
|
|
// ret.back().push_front(&c);
|
|
// }
|
|
// }
|
|
// return ret;
|
|
// }
|
|
template<typename Container>
|
|
void get_nodes_in_dir(const branch b, int dir,
|
|
std::queue<branch,Container> &nodes) const {
|
|
nodes.push_back(b);
|
|
const unsigned int dim_dir = std::abs(dir) - 1;
|
|
if (!b.front().node->children)
|
|
return;
|
|
for (int i = 0; i < 1 << dim; i++) {
|
|
if (coord{i}[dim_dir] == dir > 0)
|
|
continue;
|
|
branch child_branch(b.begin(), b.end());
|
|
child_branch.push_front({b.front().node->children->at(i),
|
|
b.front().depth - 1, b.front().width / 2,
|
|
coord{i}});
|
|
get_nodes_in_dir(child_branch, dir, nodes);
|
|
}
|
|
}
|
|
|
|
template<typename Container>
|
|
void get_adjacent_nodes(branch b, int dir,
|
|
std::queue<branch,Container> &ad_nodes) const {
|
|
auto n = b.first();
|
|
const unsigned int dim_dir = std::abs(dir) - 1;
|
|
int up = 1;
|
|
auto b_it = b.begin();
|
|
if (std::next(b_it) == b.end())
|
|
return;
|
|
for (; b_it.crd[dim_dir] == dir > 0; std::advance(b_it, 1)) {
|
|
if (std::next(b_it) == b.end())
|
|
return;
|
|
up++;
|
|
}
|
|
branch c(b.begin(), std::next(b_it));
|
|
branch p(std::next(b_it), b.end()); // TODO remove from b not cpy
|
|
c.reverse();
|
|
for (int i = 0; i < up; i += 1) {
|
|
coord pc = coord(c.front().crd).flip(dim_dir);
|
|
p.push_front({c.front().node->children->at(pc), c.front().depth - i,
|
|
c.front().width / (1 << i), pc});
|
|
c.pop_front();
|
|
}
|
|
get_nodes_in_dir(p, -dir, ad_nodes);
|
|
}
|
|
|
|
void datas_by_distance(const Point &p, cdata_access_f f) const {
|
|
branch b = get_branch_to_leaf(p);
|
|
std::deque<branch> to_visit{b};
|
|
std::unordered_set<const Node *> visited;
|
|
while (!to_visit.empty()) {
|
|
const branch& curr_visit = to_visit.front();
|
|
for (const auto &n : curr_visit) {
|
|
if (visited.contains(n.node))
|
|
break;
|
|
bool ret;
|
|
if (n.node->data_acces_ret(f, ret) && !ret)
|
|
return;
|
|
visited.insert(n.node);
|
|
}
|
|
to_visit.pop_front();
|
|
for (int i = 0; i < dim; i += 1) {
|
|
get_adjacent_nodes(curr_visit,dim+1,to_visit);//to_visit.push_back({p1, get_branch_to_leaf(p1)});
|
|
get_adjacent_nodes(curr_visit,-dim-1,to_visit);//to_visit.push_back({p1, get_branch_to_leaf(p1)});
|
|
}
|
|
}
|
|
}
|
|
// void datas_by_distance(const Point &p, cdata_access_f f) const {
|
|
// branch b = get_branch_to_leaf(p);
|
|
// std::deque<std::pair<Point, branch>> to_visit{{p, b}};
|
|
// std::unordered_set<Node *> visited;
|
|
// while (!to_visit.empty()) {
|
|
// for (auto &n : to_visit.second.front()) {
|
|
// if (visited.contains(n))
|
|
// break;
|
|
// bool ret;
|
|
// if (n.node->data_acces_ret(f, ret) && !ret)
|
|
// return;
|
|
// visited.insert(n.node);
|
|
// }
|
|
// auto search_p = to_visit.front().first;
|
|
// to_visit.pop_front();
|
|
// for (int i = 0; i < dim; i += 1) {
|
|
// Point offset = Point::Zero();
|
|
// offset(i) = b.width;
|
|
// auto p1 = search_p + offset;
|
|
// auto p2 = search_p + offset;
|
|
// to_visit.push_back({p1, get_branch_to_leaf(p1)});
|
|
// to_visit.push_back({p2, get_branch_to_leaf(p2)});
|
|
// }
|
|
// }
|
|
// }
|
|
|
|
void data_in_bb(const BBox &bb, cdata_access_f f) const {
|
|
number curr_width = root_width;
|
|
std::vector<const Node *> to_visit{root.get()};
|
|
std::vector<const Node *> to_visit_children;
|
|
bool stop = false;
|
|
while (!to_visit.empty() && !stop) {
|
|
auto cbb = to_visit.back()->bb(curr_width);
|
|
auto bbrel = cbb.relation(bb);
|
|
// std::cout << to_visit.back()->center.transpose() << " - "
|
|
// << cbb.min.transpose() << " - " << bbrel << std::endl;
|
|
bool ret;
|
|
switch (bbrel) {
|
|
case (Orthtree::BBox::CONTAINED):
|
|
to_visit.back()->all_data(f);
|
|
break;
|
|
case (Orthtree::BBox::OVERLAP):
|
|
case Orthtree::BBox::CONTAINS:
|
|
to_visit.back()->access_data_ret(f, ret);
|
|
if (to_visit.back()->access_data_ret(f, ret) && !ret)
|
|
stop = true;
|
|
if (!to_visit.back()->is_leaf()) {
|
|
// std::transform(std::execution::par_unseq,to_visit.back()->children->cbegin(),
|
|
// to_visit.back()->children->cend(),
|
|
// to_visit_children.begin(), std::addressof<const
|
|
// Node>);
|
|
for (const auto &n : *to_visit.back()->children) {
|
|
to_visit_children.push_back(&n);
|
|
}
|
|
// std::cout << "leaf "<<to_visit_children.size() << std::endl;
|
|
}
|
|
break;
|
|
default:
|
|
break;
|
|
}
|
|
to_visit.pop_back();
|
|
if (to_visit.empty()) {
|
|
to_visit.swap(to_visit_children);
|
|
to_visit_children.clear();
|
|
curr_width /= 2;
|
|
}
|
|
}
|
|
}
|
|
};
|