【HDOJ】4347 The Closest M Points

居然是KD解。

/* 4347 */
#include <iostream>
#include <sstream>
#include <string>
#include <map>
#include <queue>
#include <set>
#include <stack>
#include <vector>
#include <deque>
#include <algorithm>
#include <cstdio>
#include <cmath>
#include <ctime>
#include <cstring>
#include <climits>
#include <cctype>
#include <cassert>
#include <functional>
#include <iterator>
#include <iomanip>
using namespace std;
//#pragma comment(linker,"/STACK:102400000,1024000")

#define sti                set<int>
#define stpii            set<pair<int, int> >
#define mpii            map<int,int>
#define vi                vector<int>
#define pii                pair<int,int>
#define vpii            vector<pair<int,int> >
#define rep(i, a, n)     for (int i=a;i<n;++i)
#define per(i, a, n)     for (int i=n-1;i>=a;--i)
#define clr                clear
#define pb                 push_back
#define mp                 make_pair
#define fir                first
#define sec                second
#define all(x)             (x).begin(),(x).end()
#define SZ(x)             ((int)(x).size())
#define lson            l, mid, rt<<1
#define rson            mid+1, r, rt<<1|1

int id, n, m, K;

typedef struct node_t {
    int x[5];

    friend bool operator< (const node_t& a, const node_t& b) {
        return a.x[id] < b.x[id];
    }

    double Distance(const node_t& a) {
        __int64 ret = 0;

        rep(i, 0, K)
            ret += 1LL*(a.x[i]-x[i])*(a.x[i]-x[i]);

        return ret;
    }

    void print() {
        printf("%d", x[0]);
        rep(i, 1, K)
            printf(" %d", x[i]);
        putchar('\n');
    }

} node_t;

typedef struct Node {
    __int64 d;
    node_t p;

    Node() {}
    Node(__int64 d, node_t& p):
        d(d), p(p) {}

    friend bool operator< (const Node& a, const Node& b) {
        return a.d < b.d;
    }

} Node;

const int maxn = 50005;
node_t nd[maxn];
vector<node_t> ans;

typedef struct KD_tree{
    static const int maxd = 5;
    static const int maxn = 50005;
    node_t P[maxn<<2];
    int idx[maxn<<2];
    priority_queue<Node> Q;

    void Build(int deep, int l, int r, int rt) {
        if (l > r)
            return ;
        idx[rt] = deep % K;
        if (l == r) {
            P[rt] = nd[l];
            return ;
        }

        id = idx[rt];
        int mid = (l + r) >> 1;
        nth_element(nd+l, nd+mid, nd+r+1);
        P[rt] = nd[mid];
        Build(deep+1, l, mid-1, rt<<1);
        Build(deep+1, mid+1, r, rt<<1|1);
    }

    void Query(node_t x, int l, int r, int rt) {
        if (l > r)
            return ;
        
        int id = idx[rt];
        __int64 tmp = x.Distance(P[rt]);
        if (l == r) {
            if (SZ(Q) < m) {
                Q.push(Node(tmp, P[rt]));
            } else {
                if (tmp < Q.top().d) {
                    Q.pop();
                    Q.push(Node(tmp, P[rt]));
                }
            }
            return ;
        }
        
        int mid = (l + r) >> 1;

        if (P[rt].x[id] >= x.x[id]) {
            Query(x, l, mid-1, rt<<1);
            if (SZ(Q) < m) {
                Q.push(Node(tmp, P[rt]));
                Query(x, mid+1, r, rt<<1|1);
            } else {
                if (tmp < Q.top().d) {
                    Q.pop();
                    Q.push(Node(tmp, P[rt]));
                }
                if (1LL*(x.x[id]-P[rt].x[id])*(x.x[id]-P[rt].x[id]) < Q.top().d) {
                    Query(x, mid+1, r, rt<<1|1);
                }
            }
        } else {
            Query(x, mid+1, r, rt<<1|1);
            if (SZ(Q) < m) {
                Q.push(Node(tmp, P[rt]));
                Query(x, l, mid-1, rt<<1);
            } else {
                if (tmp < Q.top().d) {
                    Q.pop();
                    Q.push(Node(tmp, P[rt]));
                }
                if (1LL*(x.x[id]-P[rt].x[id])*(x.x[id]-P[rt].x[id]) < Q.top().d) {
                    Query(x, l, mid-1, rt<<1);
                }
            }
        }
    }

    void Dump() {
        ans.clr();
        while (!Q.empty()) {
            ans.pb(Q.top().p);
            Q.pop();
        }
    }
} KD_tree;

KD_tree kd;

void solve() {
    int q, sz;
    node_t p;

    kd.Build(0, 0, n-1, 1);
    scanf("%d", &q);
    while (q--) {
        rep(j, 0, K)
            scanf("%d", &p.x[j]);
        scanf("%d", &m);
        kd.Query(p, 0, n-1, 1);
        kd.Dump();
        sz = SZ(ans);
        printf("the closest %d points are:\n", m);
        per(i, 0, sz) {
            ans[i].print();
        }
    }
}

int main() {
    ios::sync_with_stdio(false);
    #ifndef ONLINE_JUDGE
        freopen("data.in", "r", stdin);
        freopen("data.out", "w", stdout);
    #endif

    while (scanf("%d %d", &n, &K)!=EOF) {
        rep(i, 0, n)
            rep(j, 0, K)
                scanf("%d", &nd[i].x[j]);
        solve();
    }

    #ifndef ONLINE_JUDGE
        printf("time = %d.\n", (int)clock());
    #endif

    return 0;
}

数据发生器。

from copy import deepcopy
from random import randint, shuffle
import shutil
import string


def GenDataIn():
    with open("data.in", "w") as fout:
        t = 10
        bound = 10**5
        for tt in xrange(t):
            n = randint(100, 200)
            k = randint(1, 5)
            fout.write("%d %d\n" % (n, k))
            for i in xrange(n):
                L = []
                for j in xrange(k):
                    x = randint(-1000, 1000)
                    L.append(x)
                fout.write(" ".join(map(str, L)) + "\n")
            q = randint(20, 30)
            fout.write("%d\n" % (q))
            for qq in xrange(q):
                L = []
                for j in xrange(k):
                    x = randint(-1000, 1000)
                    L.append(x)
                fout.write(" ".join(map(str, L)) + "\n")
                m = randint(1, 10)
                fout.write("%d\n" % (m))
            
                
def MovDataIn():
    desFileName = "F:\eclipse_prj\workspace\hdoj\data.in"
    shutil.copyfile("data.in", desFileName)

    
if __name__ == "__main__":
    GenDataIn()
    MovDataIn()