显示代码纯文本
#include <iostream>
#include <string.h>
#include <algorithm>
#include <stdio.h>
#include <math.h>
#include <queue>
typedef long long TYPE;
using namespace std;
#define N 50005
#define lson rt << 1
#define rson rt << 1 | 1
#define Pair pair<double, Node>
#define Sqrt2(x) ( (x) * (x) * (x) *(x) )
FILE *IN,*OUT;
TYPE n, k, idx;
struct Node
{
TYPE feature[5]; //定义属性数组
bool operator < (const Node &u) const
{
return feature[idx] < u.feature[idx];
}
}_data[N]; //_data[]数组代表输入的数据
priority_queue<Pair> Q; //队列Q用于存放离p最近的m个数据
class KDTree{
public:
void Build(TYPE, TYPE, TYPE, TYPE); //建树
void Query(Node, TYPE, TYPE, TYPE); //查询
private:
Node data[4 * N]; //data[]数组代表K-D树的所有节点数据
TYPE flag[4 * N]; //用于标记某个节点是否存在，1表示存在，-1表示不存在
}kd;
//建树步骤，参数dept代表树的深度
void KDTree::Build(TYPE l, TYPE r, TYPE rt, TYPE dept)
{
if(l > r) return;
flag[rt] = 1; //表示编号为rt的节点存在
flag[lson] = flag[rson] = -1; //当前节点的孩子暂时标记不存在
idx = dept % k; //按照编号为idx的属性进行划分
TYPE mid = (l + r) >> 1;
nth_element(_data + l, _data + mid, _data + r + 1); //nth_element()为STL中的函数
data[rt] = _data[mid];
Build(l, mid - 1, lson, dept + 1); //递归左子树
Build(mid + 1, r, rson, dept + 1); //递归右子树
}
//查询函数，寻找离p最近的m个特征属性
void KDTree::Query(Node p, TYPE m, TYPE rt, TYPE dept)
{
if(flag[rt] == -1) return; //不存在的节点不遍历
Pair cur(0, data[rt]); //获取当前节点的数据和到p的距离
for(TYPE i = 0; i < k; i++)
cur.first += Sqrt2(cur.second.feature[i] - p.feature[i]);
TYPE dim = dept % k; //跟建树一样，这样能保证相同节点的dim值不变
bool fg = 0; //用于标记是否需要遍历右子树
TYPE x = lson;
TYPE y = rson;
if(p.feature[dim] >= data[rt].feature[dim]) //数据p的第dim个特征值大于等于当前的数据，则需要进入右子树
swap(x, y);
if(~flag[x]) Query(p, m, x, dept + 1); //如果节点x存在，则进入子树继续遍历
//以下是回溯过程，维护一个优先队列
if(Q.size() < m) //如果队列没有满，则继续放入
{
Q.push(cur);
fg = 1;
}
else
{
if(cur.first < Q.top().first) //如果找到更小的距离，则用于替换队列Q中最大的距离的数据
{
Q.pop();
Q.push(cur);
}
if(Sqrt2(p.feature[dim] - data[rt].feature[dim]) < Q.top().first)
{
fg = 1;
}
}
if(~flag[y] && fg)
Query(p, m, y, dept + 1);
}
//输出结果
void Print(Node data,Node p)
{
TYPE ans = 0;
for(TYPE i = 0; i < k; i++)
ans += Sqrt2(data.feature[i] - p.feature[i]);
fprintf(OUT,"%lld\n", ans);
}
int main()
{
IN=fopen("Keller_Deal.in","r");
OUT=fopen("Keller_Deal.out","w");
TYPE t,m;k=5;
while(fscanf(IN,"%lld%lld", &n, &t)!=EOF)
{
for(TYPE i = 0; i < n; i++){
TYPE ds=0;
fscanf(IN,"%lld",&ds);
for(TYPE j = 0; j < ds; j++)
fscanf(IN,"%lld", &_data[i].feature[j]);
for(TYPE j =ds; j < 5; j++)
_data[i].feature[j] = 0;
}
kd.Build(0, n - 1, 1, 0);
while(t--)
{
Node p;TYPE kss=0;
fscanf(IN,"%lld%lld",&kss, &m);
for(TYPE i = 0; i < kss; i++)
fscanf(IN,"%lld", &p.feature[i]);
for(TYPE i = kss; i < 5; i++)
p.feature[i] = 0;
while(!Q.empty()) Q.pop(); //事先需要清空优先队列
kd.Query(p, m, 1, 0);
fprintf(OUT,"the closest %lld strings' distance are:\n", m);
Node tmp[25];
for(TYPE i = 0; !Q.empty(); i++)
{
tmp[i] = Q.top().second;
Q.pop();
}
for(TYPE i = m - 1; i >= 0; i--)
Print(tmp[i],p);
}
}
return 0;
}
YGOI_真神名曰驴蛋蛋