K Closest Points to Origin LeetCode solution in Python

1
# Runtime: 2003 ms (Top 6.82%) | Memory: 20.5 MB (Top 21.33%)
2

3

4
def cal(ele):
5
    return ele[0] ** 2 + ele[1] ** 2
6

7

8
def partition(arr, start, end):
9
    # We must take a pivot element randomly to make Quick select work faster
10
    rdIdx = randint(start, end)
11
    arr[rdIdx], arr[end] = arr[end], arr[rdIdx]
12
    pivot = arr[end]
13
    i = start - 1
14
    pivot_dis = cal(pivot)
15
    for j in range(start, end):
16
        if cal(arr[j]) <= pivot_dis:
17
            i += 1
18
            arr[j], arr[i] = arr[i], arr[j]
19

20
    arr[i + 1], arr[end] = arr[end], arr[i + 1]
21
    return i + 1
22

23

24
def qSelect(arr, kth, start, end):
25
    if start < end:
26
        pv = partition(arr, start, end)
27
        # _______________________
28
        # | Left |ele| Right|
29
        # ------------------------
30
        # pv
31
        # after partition function call, pv is the index that sacrify:
32
        # all elements in Left will smaller than ele
33
        # all elements in Right side will greater than ele
34
        if pv == kth:  #
35
            return
36
        if kth < pv:
37
            return qSelect(arr, kth, start, pv - 1)
38
        else:
39
            return qSelect(arr, kth, pv + 1, end)
40

41

42
# Space O (logn) because of using recursion
43
# Time: Average case: O(N)
44
# Worst case: O(N**2)
45
class Solution:
46
    def kClosest(self, points, k):
47
        # print(points)
48
        qSelect(
49
            points, k - 1, 0, len(points) - 1
50
        )  # kth smallest number will be at (k-1) index in sorted array
51
        return points[:k]