HYSBZ bzoj 4520 K远点对

Problem

www.lydsy.com/JudgeOnline/problem.php?id=4520
vjudge.net/contest/187908#problem

Reference

BZOJ4520 [Cqoi2016]K远点对
【kd-tree】BZOJ4520 CQOI2016K远点对

Meaning

给出平面上 n 个点，问第距离 k 远的点对的距离的平方是多少

Analysis

暴力的想法，对每个点查询其它点与它的距离，找第 2k 个（因为每个距离都被算了两次）。用 K-D树来优化搜索（剪枝？），维护一个记录点距的、大小为 2k 的小根堆。

Notes

学了一种新的 K-D树写法，感觉更正宗？（更像用指针动态分配的那种，只是换成数组模拟指针）上一篇的那种写法空间复杂度不会算，数组大小开得很迷…
代码中的evaluate()是用来估价的，判断是否有可能搜到更远的点距，决定是否要往下搜，也算一种剪枝吧。估价时用到当前结点分割的那两个子区域里的最大和最小的横、纵坐标，在建树时处理出来。

Code

#include <algorithm>
#include <iostream>
#include <queue>
#include <utility>
#include <vector>
using namespace std;
typedef long long LL;
const int N = 100000, K = 100, D = 2;

int idx;

struct node
{
    int v[D];
    bool operator < (const node &rhs) const
    {
        return v[idx] < rhs.v[idx];
    }
} kdt[N], tar;

int big[N][D]; // big[i][j]：结点 i 的区域里第 j 维的最大值
int sml[N][D]; // sml[i][j]：结点 i 的区域里第 j 维的最小值
int lc[N]; // 左子树根的下标
int rc[N]; // 右子树根的下标

void pushup(int rt)
{
    for(int i = 0; i < D; ++i)
    {
        int &b = big[rt][i], &s = sml[rt][i];
        if(~lc[rt]) // 存在左子树
        {
            b = max(b, big[lc[rt]][i]);
            s = min(s, sml[lc[rt]][i]);
        }
        if(~rc[rt]) // 存在右子树
        {
            b = max(b, big[rc[rt]][i]);
            s = min(s, sml[rc[rt]][i]);
        }
    }
}

int build(int l, int r, int d)
{
    // 当前考虑第 idx 维
    idx = d & 1;
    int m = l + r >> 1;
    // 左右子树标记为不存在先
    lc[m] = rc[m] = -1;
    nth_element(kdt + l, kdt + m, kdt + r + 1);
    // 初始化 big 和 sml
    for(int i = 0; i < D; ++i)
        big[m][i] = sml[m][i] = kdt[m].v[i];
    if(l < m) // 左子树存在
        lc[m] = build(l, m - 1, d + 1);
    if(r > m) // 右子树存在
        rc[m] = build(m + 1, r, d + 1);
    // 更新 big 和 sml
    pushup(m);
    return m;
}

long long sqr(long long x)
{
    return x * x;
}

long long evaluate(int id)
{
    long long res = 0;
    for(int i = 0; i < D; ++i)
        res += max( // 找尽量大的距离来估价
            sqr(tar.v[i] - big[id][i]),
            sqr(tar.v[i] - sml[id][i])
        );
    return res;
}

// greater -> 小顶
priority_queue<LL, vector<LL>, greater<LL> > que;

void query(int now)
{
    long long d = 0;
    // 当前点和目标点的距离
    for(int i = 0; i < D; ++i)
        d += sqr(tar.v[i] - kdt[now].v[i]);
    // 更新堆顶
    if(d > que.top())
    {
        que.pop();
        que.push(d);
    }

    long long dl = 0, dr = 0;
    if(~lc[now]) // 估价左子树
        dl = evaluate(lc[now]);
    if(~rc[now]) // 估价右子树
        dr = evaluate(rc[now]);

    if(dl > dr) // 左子树估价更高 -> 先搜左边
    {
        if(dl > que.top())
            query(lc[now]);
        if(dr > que.top())
            query(rc[now]);
    }
    else // 先搜右边
    {
        if(dr > que.top())
            query(rc[now]);
        if(dl > que.top())
            query(lc[now]);
    }
}

int main()
{
    ios::sync_with_stdio(false);
    cin.tie(0);
    int n, k;
    cin >> n >> k;
    for(int i = 0; i < n; ++i)
        cin >> kdt[i].v[0] >> kdt[i].v[1];

    int rt = build(0, n - 1, 0);
    // 大小为 2k 的堆
    for(int i = 0; i < k << 1; ++i)
        que.push(0LL);
    for(int i = 0; i < n; ++i)
    {
        tar = kdt[i];
        query(rt);
    }

    cout << que.top() << '\n';
    while(!que.empty())
        que.pop();
    return 0;
}