| 1 | #define INC(c) ((c) == K - 1 ? 0 : (c) + 1) |
| 2 | template <int K> struct kd_tree { |
| 3 | struct pt { |
| 4 | double coord[K]; |
| 5 | pt() {} |
| 6 | pt(double c[K]) { rep(i,0,K) coord[i] = c[i]; } |
| 7 | double dist(const pt &other) const { |
| 8 | double sum = 0.0; |
| 9 | rep(i,0,K) sum += pow(coord[i] - other.coord[i], 2.0); |
| 10 | return sqrt(sum); } }; |
| 11 | struct cmp { |
| 12 | int c; |
| 13 | cmp(int _c) : c(_c) {} |
| 14 | bool operator ()(const pt &a, const pt &b) const { |
| 15 | for (int i = 0, cc; i <= K; i++) { |
| 16 | cc = i == 0 ? c : i - 1; |
| 17 | if (abs(a.coord[cc] - b.coord[cc]) > EPS) |
| 18 | return a.coord[cc] < b.coord[cc]; |
| 19 | } |
| 20 | return false; } }; |
| 21 | struct bb { |
| 22 | pt from, to; |
| 23 | bb(pt _from, pt _to) : from(_from), to(_to) {} |
| 24 | double dist(const pt &p) { |
| 25 | double sum = 0.0; |
| 26 | rep(i,0,K) { |
| 27 | if (p.coord[i] < from.coord[i]) |
| 28 | sum += pow(from.coord[i] - p.coord[i], 2.0); |
| 29 | else if (p.coord[i] > to.coord[i]) |
| 30 | sum += pow(p.coord[i] - to.coord[i], 2.0); |
| 31 | } |
| 32 | return sqrt(sum); } |
| 33 | bb bound(double l, int c, bool left) { |
| 34 | pt nf(from.coord), nt(to.coord); |
| 35 | if (left) nt.coord[c] = min(nt.coord[c], l); |
| 36 | else nf.coord[c] = max(nf.coord[c], l); |
| 37 | return bb(nf, nt); } }; |
| 38 | struct node { |
| 39 | pt p; node *l, *r; |
| 40 | node(pt _p, node *_l, node *_r) |
| 41 | : p(_p), l(_l), r(_r) { } }; |
| 42 | node *root; |
| 43 | // kd_tree() : root(NULL) { } |
| 44 | kd_tree(vector<pt> pts) { |
| 45 | root = construct(pts, 0, (int)size(pts) - 1, 0); } |
| 46 | node* construct(vector<pt> &pts, int from, int to, int c) { |
| 47 | if (from > to) return NULL; |
| 48 | int mid = from + (to - from) / 2; |
| 49 | nth_element(pts.begin() + from, pts.begin() + mid, |
| 50 | pts.begin() + to + 1, cmp(c)); |
| 51 | return new node(pts[mid], |
| 52 | construct(pts, from, mid - 1, INC(c)), |
| 53 | construct(pts, mid + 1, to, INC(c))); } |
| 54 | bool contains(const pt &p) { return _con(p, root, 0); } |
| 55 | bool _con(const pt &p, node *n, int c) { |
| 56 | if (!n) return false; |
| 57 | if (cmp(c)(p, n->p)) return _con(p, n->l, INC(c)); |
| 58 | if (cmp(c)(n->p, p)) return _con(p, n->r, INC(c)); |
| 59 | return true; } |
| 60 | void insert(const pt &p) { _ins(p, root, 0); } |
nothing calls this directly
no outgoing calls
no test coverage detected