1
2
3
4
5
6
7

Input: points = [[1,3],[-2,2]], K = 1
Output: [[-2,2]]
Explanation: 
The distance between (1, 3) and the origin is sqrt(10).
The distance between (-2, 2) and the origin is sqrt(8).
Since sqrt(8) < sqrt(10), (-2, 2) is closer to the origin.
We only want the closest K = 1 points from the origin, so the answer is just [[-2,2]].

1
2
3

Input: points = [[3,3],[5,-1],[-2,4]], K = 2
Output: [[3,3],[-2,4]]
(The answer [[-2,4],[3,3]] would also be accepted.)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54

func kClosest(points [][]int, K int) [][]int {
    res := make(PQ, len(points))
    
    for idx:=range points{
        res[idx] = getn(points[idx])
    }

    heap.Init(&res)
    
    resa := make([][]int, 0)
    
    for i:=0;i<K;i++ {
        r := heap.Pop(&res)
        resa = append(resa, (r).(*N).A )
    }
    return resa
}
type PQ []*N
type N struct {
    A []int
    D float64
}

func (v PQ) Len() int{
    return len(v)
}

func (v PQ) Less(a,b int) bool{
    return v[a].D<v[b].D
}

func (v *PQ) Pop() interface{}{
    old := *v
    n := len(old)
    item := old[n-1]
    *v = old[:n-1]

    return item
}

func (v *PQ) Push(el interface{}) {
    *v = append(*v, el.(*N))
}

func (v PQ) Swap(a,b int) {
    v[a],v[b] = v[b],v[a]
}

func getn(a []int) *N{
    return &N{
        A: a,
        D: math.Sqrt(float64(a[0])*float64(a[0]) +float64(a[1])*float64(a[1])),
    }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52

func kClosest(points [][]int, K int) [][]int {
    f := sort(points)
    return f[:K]
}

func sort(points [][]int) [][]int{
    if len(points) == 1 {
        return points
    }
    if  len(points) ==2 {
        
        if dist(points[0])>dist(points[1]){ 
            points[0],points[1] = points[1],points[0]
        }
        
        return points
    }
    
    m:= len(points)/2
    left:=sort(points[0: m])
    right:=sort(points[m:len(points)])
    res := make([][]int,0,len(points))
    i:=0
    j:=0
    for i<len(left) || j<len(right) {
        if i<len(left) && j<len(right) {
            if dist(left[i])<=dist(right[j]) {
                res = append(res, left[i])
                i++
            }else {
                res = append(res, right[j])
                j++
            }
            continue
        }
        if i<len(left) {
            res = append(res, left[i])
            i++
            continue
        }
        if j<len(right) {
            res = append(res, right[j])
            j++
            continue
        }
    }
    return res
}

func dist(p []int) int {
    return p[0]*p[0]+p[1]*p[1]
}