hdu 4347

#include<bits/stdc++.h>
#define sq(x) (x)*(x)
#define N (55555)
 
using namespace std;
 
int idx,k,n,m,q;
 
struct Node
{
    int x[5];
    bool operator < (const Node &u) const
    {
        return x[idx] < u.x[idx];
    }
} P[N];

typedef pair<double,Node> PDN;
priority_queue<PDN> que;    
 
struct KD_Tree
{
    int sz[N<<2]; Node p[N<<2];
 
    void nthelement(int s,int k,int e,int idx)
    {
        if (s >= e)
        {
            return;
        }
        //int mid = random.Next(s, e);
        // int mid = (s + e) / 2;
        // Node tmp = P[mid];
        // P[mid] = P[s];
        // P[s] = tmp;
        Node md = P[s];
        int l = 0;
        for (int i = s + 1; i < e; i++)
        {
            if (P[i].x[idx] <= md.x[idx])
            {
                ++l;
                Node tmp = P[s + l];
                P[s + l] = P[i];
                P[i] = tmp;
            }
        }

        P[s] = P[l + s];
        P[l + s] = md;
        if (l + 1 == k)
            return;
        if (l + 1 < k)
        {
            //ss.Push(new F2(s + l + 1, k - l - 1, e));
            nthelement(s + l + 1, k - l - 1, e, idx);
        }
        else
        {
            //ss.Push(new F2(s, l + 1 - k, s + l + 1));
            nthelement(s, k, s + l, idx);
        }
    }

    void build(int i,int l,int r,int dep)
    {
        if (l>r) return;
        int mid=(l+r)>>1;
        idx=dep%k;sz[i]=r-l;
        sz[i<<1]=sz[i<<1|1]=-1;
        //nth_element(P+l,P+mid,P+r+1);
        nthelement(l,mid+1-l,r+1,idx);
        p[i]=P[mid];
        build(i<<1,l,mid-1,dep+1);
        build(i<<1|1,mid+1,r,dep+1);
    }
 
    void query(int i,int m,int dep,Node a)
    {
        if (sz[i]==-1) return;
        PDN tmp=PDN(0,p[i]);
        for(int j=0;j<k;j++)
            tmp.first+=sq(tmp.second.x[j]-a.x[j]);
        int lc=i<<1,rc=i<<1|1,dim=dep%k,flag=0;
        if (a.x[dim]>=p[i].x[dim]) swap(lc,rc);
        if (~sz[lc]) query(lc,m,dep+1,a);
        if (que.size()<m) que.push(tmp),flag=1;
        else
        {
            if (tmp.first<que.top().first) que.pop(),que.push(tmp);
            if (sq(a.x[dim]-p[i].x[dim])<que.top().first) flag=1;
        }
        if (~sz[rc]&&flag) query(rc,m,dep+1,a);
    }
} KDT;
 
int main()
{
    while(~scanf("%d%d",&n,&k))
    {
        for(int i=0;i<n;i++)
            for(int j=0;j<k;j++)
                scanf("%d",&P[i].x[j]);
        
        KDT.build(1,0,n-1,0);
        scanf("%d",&q);
        while(q--)
        {
            Node now;
            for(int i=0;i<k;i++)
                scanf("%d",&now.x[i]);
            scanf("%d",&m); int t=0;
            KDT.query(1,m,0,now); Node pp[21];
            for(;!que.empty();que.pop())
                pp[++t]=que.top().second;
            printf("the closest %d points are:
",t);
            for(int i=m;i>0;i--)
            {
                printf("%d",pp[i].x[0]);
                for(int j=1;j<k;j++)
                    printf(" %d",pp[i].x[j]);
                puts("");
            }
        }
    }
 
    return 0;
 
}

C# TLE

using System;
using System.Collections.Generic;
using System.Data.SqlTypes;
using System.Linq;

namespace kdtree
{
    public class F
    {
        public F(int i,int cnt,int k,Node nd)
        {
            this.i = i;
            this.cnt = cnt;
            this.k = k;
            this.nd = nd;
        }
        public int i, cnt, k;
        public Node nd;
    }

    public class F2
    {
        public F2(int s, int k, int e)
        {
            this.s = s;
            this.k = k;
            this.e = e;
        }
        public int s, k, e;
    }
    public class Node:IComparable<Node>
    {
        public List<int> a;
        public int u;
        public int d;

        public int CompareTo(Node obj)
        {
            var val = d.CompareTo(obj.d);
            if (val != 0)
                return val;
            return u.CompareTo(obj.u);
        }
    }

    public class treekd
    {
        public int dim;
        public void init(int n,int dim,List<Node> points)
        {
            int N = n << 2;
            sz = new int[N];
            p = new Node[N];
            this.dim = dim;
            this.points = points;
        }
        int[] sz;
        List<Node> points;
        Node[] p;
        public SortedSet<Node> que;
        public void initquery()
        {
            que = new SortedSet<Node>();
        }
        
        public void createnode()
        {

        }
        public void update()
        {

        }
        public int sq(int a) { return a * a; }
        public void swap(ref int a,ref int b)
        {
            a = a + b;
            b = a - b;
            a = a - b;
        }
        public void query(int i,int cnt,int k,Node nd)
        {
            Stack<F> s = new Stack<F>();
            s.Push(new F(i,cnt,k,nd));
            while (s.Count > 0)
            {
                F f = s.Pop();
                i = f.i;cnt = f.cnt;k = f.k;nd = f.nd;
                if (sz[i] == -1) continue;
                Node tmp = new Node() { u = i, a = p[i].a };
                for (int ii = 0; ii < dim; ii++)
                {
                    tmp.d += sq(tmp.a[ii] - nd.a[ii]);
                }
                int lc = i << 1, rc = i << 1 | 1;
                int ndim = k % dim;
                bool flag = false;
                if (nd.a[ndim] >= p[i].a[ndim]) swap(ref lc, ref rc);
                if (sz[lc] != -1)
                {
                    s.Push(new F(lc,cnt,k+1,nd));
                    //query(lc, cnt, k + 1, nd);
                }
                if (que.Count < cnt)
                {
                    que.Add(tmp);
                    flag = true;
                }
                else
                {
                    if (tmp.d < que.Last().d)
                    {
                        que.Remove(que.Last());
                        que.Add(tmp);
                    }
                    if (sq(nd.a[ndim] - p[i].a[ndim]) < que.Last().d) flag = true;
                }
                if (sz[rc] != -1 && flag == true)
                {
                    s.Push(new F(rc, cnt, k + 1, nd));
                    //query(rc, cnt, k + 1, nd);
                }
            }
        }
        static Random random = new Random();
        public void nth_element(ref List<Node> list,int s,int k,int e,int idx)
        {
            Stack<F2> ss = new Stack<F2>();
            ss.Push(new F2(s, k, e));
            while (ss.Count > 0)
            {
                F2 f = ss.Pop();
                s = f.s; k = f.k; e = f.e;
                if (s >= e)
                {
                    continue;
                }
                //int mid = random.Next(s, e);
                //int mid = (s + e) / 2;
                //Node tmp = list[mid];
                //list[mid] = list[s];
                //list[s] = tmp;
                Node md = list[s];
                int l = 0;
                for (int i = s + 1; i < e; i++)
                {
                    if (list[i].a[idx] <= md.a[idx])
                    {
                        ++l;
                        Node tmp = list[s + l];
                        list[s + l] = list[i];
                        list[i] = tmp;
                    }
                }

                list[s] = list[l + s];
                list[l + s] = md;
                if (l + 1 == k)
                    return;
                if (l + 1 < k)
                {
                    ss.Push(new F2(s + l + 1, k - l - 1, e));
                    //nth_element(ref list, s + l + 1, k - l - 1, e);
                }
                else
                {
                    ss.Push(new F2(s, k, s + l));
                    //nth_element(ref list, s, l + 1 - k, s + l + 1);
                }
            }
        }
        public void build(int i, int l,int r,int k)
        {
            if (l > r) return;
            int mid = (l + r) >> 1;
            int idx = k % dim;
            sz[i] = r - l;
            sz[i << 1] = sz[i << 1 | 1] = -1;
            nth_element(ref points, l, mid+1-l, r+1, idx);
            //var list=points.Skip(l).Take(r-l+1).OrderBy(q => q.a[idx]).ToList();
            //for (int ii = l; ii <= r; ii++)
            //    points[ii] = list[ii-l];
            p[i] = points[mid];
            build(i << 1, l, mid - 1, k + 1);
            build(i << 1 | 1, mid + 1, r, k + 1);
        }
    }
    class Program
    {
        static public List<int> GetLine()
        {
            var input = Console.ReadLine();
            if (input == null)
                return null;
            return input.Split(' ').Select(p=>int.Parse(p)).ToList();
        }
        static public int GetInt()
        {
            return Console.ReadLine().Split(' ').Select(p => int.Parse(p)).ToList()[0];
        }
        static public void nth_element(ref List<int> list, int s, int k, int e)
        {
            if (s >= e|| k==0)
            {
                return;
            }
            //int mid = (s + e) / 2;
            //int tmp = list[mid];
            //list[mid] = list[s];
            //list[s] = tmp;
            int md = list[s];
            int l = 0;
            for (int i = s+1; i < e; i++)
            {
                if (list[i] <= md)
                {
                    ++l;
                    var tmp = list[s+l];
                    list[s+l] = list[i];
                    list[i] = tmp;
                }
            }

            list[s]=list[l+s];
            list[l+s] = md;
            if (l + 1 == k)
                return;
            if (l + 1 < k)
            {
                nth_element(ref list, s+l+1, k - (l + 1), e);
            }
            else
            {
                nth_element(ref list, s, k, s + l);
            }
        }
        static void Main(string[] args)
        {
            //List<int> list1 = new List<int>() { 8, 3, 4, 5, 1, 6, 8, 2, 3, 0 };
            //List<int> list1 = new List<int>();
            //for (int i = 0; i < 100; i++)
            //{
            //    list1.Add(new Random().Next(0, 100));
            //}
            ////List<int> list1
            //for (int i = 1; i < list1.Count; i++)
            //{
            //    var list11 = list1.ToList();
            //    nth_element(ref list11, 0, i + 1, list11.Count);
            //    Console.WriteLine(string.Join(",", list1));
            //    var list12 = list1.OrderBy(p => p).ToList();
            //    if (list11[i] != list12[i])
            //    {
            //        Console.WriteLine(list11[i]+","+list12[i]);
            //    }
            //}
            //List<int> list2 = new List<int>() { 10, 4, 6, 3, 7, 3, 7, 11, 8, 6, 10, 2 };
            //List<int> list3 = new List<int>() { 2, 7, 6, 3, 8, 3, 14, 3, 8, 6, 10, 2 };
            List<int> input;
            while ((input = GetLine()) != null&&input.Count>0)
            {
                treekd kdt = new treekd();
                var n = input[0];
                var k = input[1];
                List<Node> nodes = new List<Node>();
                for (int i = 0; i < n; i++)
                {
                    input = GetLine();
                    nodes.Add(new Node());
                    nodes[i].a = input;
                }
                kdt.init(n, k, nodes);
                kdt.build(1, 0, n - 1, 0);
                var t = GetInt();
                while (t-- > 0)
                {
                    kdt.initquery();
                    Node now = new Node();
                    now.a = GetLine();
                    int cnt = GetInt();
                    kdt.query(1, cnt, 0, now);
                    Console.WriteLine("the closest " + cnt + " points are:");
                    foreach (var item in kdt.que.ToList().Take(cnt))
                    {
                        Console.WriteLine(string.Join(" ", item.a));
                    }
                }
            }
        }
    }
}