:warning: Kd木(WIP)
(data_structure/kd_tree.hpp)

Code

#pragma once
#include<vector>

/**
 * @brief Kd木(WIP)
 */

template<typename T=long long,int sz=2>
struct kd_tree{
    typedef array<T,sz> Point;
    Point point;
    //vector<Point> v;
	long long v_size;
    int axis,depth;
    kd_tree *l=nullptr,*r=nullptr;
    kd_tree(vector<Point> v=vector<Point>{},int depth=0):v_size(v.size()),depth(depth){
        build(v,depth);
    }
    inline bool comp(const Point& s,const Point& t){
        rep(i,sz)if(s[i]>t[i])return 0;
        return 1;
    }
    inline void build(vector<Point> v,const int& depth){
        if(v_size==0)return;
        axis=depth%tuple_size<Point>::value;
        sort(all(v),[this](const auto s,const auto t){return s[this->axis]<t[this->axis];});
        point=v[v.size()/2];
        vector<Point> lv(v.size()/2),rv(v.size()-v.size()/2);
        rep(i,lv.size())lv[i]=v[i];
    	rep(i,rv.size())rv[i]=v[i+lv.size()];
        if(lv.size()&&v.size()!=1)l=new kdtree(lv,depth+1);
        if(rv.size()&&v.size()!=1)r=new kdtree(rv,depth+1);
    }
	void insert(Point p){
		v_size++;
        if(v_size==1){
            return;
        }
        if(p[axis]<point[axis]){
            if(!l)l=new kdtree(vector<Point>{p},depth+1);
            else l->insert(p);
        }else{
            if(!r)r=new kdtree(vector<Point>{p},depth+1);
            else r->insert(p);
        }
    }
	vector<Point> search(const Point& lp,const Point& rp){
        if(v_size==1&&comp(lp,point)&&comp(point,rp))return vector<Point>{point};
        std::vector<Point> res;
        if(r&&point[axis]<=rp[axis]){
            res=r->search(lp,rp);
        }
        if(l&&lp[axis]<=point[axis]){
            vector<Point> tmp=l->search(lp,rp);
			res.insert(res.end(),all(tmp));
        }
        return res;
    }
};
#line 2 "data_structure/kd_tree.hpp"
#include<vector>

/**
 * @brief Kd木(WIP)
 */

template<typename T=long long,int sz=2>
struct kd_tree{
    typedef array<T,sz> Point;
    Point point;
    //vector<Point> v;
	long long v_size;
    int axis,depth;
    kd_tree *l=nullptr,*r=nullptr;
    kd_tree(vector<Point> v=vector<Point>{},int depth=0):v_size(v.size()),depth(depth){
        build(v,depth);
    }
    inline bool comp(const Point& s,const Point& t){
        rep(i,sz)if(s[i]>t[i])return 0;
        return 1;
    }
    inline void build(vector<Point> v,const int& depth){
        if(v_size==0)return;
        axis=depth%tuple_size<Point>::value;
        sort(all(v),[this](const auto s,const auto t){return s[this->axis]<t[this->axis];});
        point=v[v.size()/2];
        vector<Point> lv(v.size()/2),rv(v.size()-v.size()/2);
        rep(i,lv.size())lv[i]=v[i];
    	rep(i,rv.size())rv[i]=v[i+lv.size()];
        if(lv.size()&&v.size()!=1)l=new kdtree(lv,depth+1);
        if(rv.size()&&v.size()!=1)r=new kdtree(rv,depth+1);
    }
	void insert(Point p){
		v_size++;
        if(v_size==1){
            return;
        }
        if(p[axis]<point[axis]){
            if(!l)l=new kdtree(vector<Point>{p},depth+1);
            else l->insert(p);
        }else{
            if(!r)r=new kdtree(vector<Point>{p},depth+1);
            else r->insert(p);
        }
    }
	vector<Point> search(const Point& lp,const Point& rp){
        if(v_size==1&&comp(lp,point)&&comp(point,rp))return vector<Point>{point};
        std::vector<Point> res;
        if(r&&point[axis]<=rp[axis]){
            res=r->search(lp,rp);
        }
        if(l&&lp[axis]<=point[axis]){
            vector<Point> tmp=l->search(lp,rp);
			res.insert(res.end(),all(tmp));
        }
        return res;
    }
};
Back to top page