1+ '''
2+ K Closest Points
3+
4+ Given an array with points, and another one point. Find the closest K points of the array to the given point.
5+
6+ Input: [(0, 1), (2, 1), (3, 3), (1, 2)], (2, 2), 2
7+ Output: [(2, 1), (1, 2)]
8+
9+ =========================================
10+ Same solution as the nth_smallest.py, in this case we're looking for the K smallest DISTANCES.
11+ Based on the quick sort algorithm (pivoting, divide and conquer).
12+ More precisly in-place quick sort (without using additional space).
13+ Time Complexity: O(N) , O(N + N/2 + N/4 + N/8 + ... + 1 = 2*N = N)
14+ Space Complexity: O(K) , length of the output array
15+ '''
16+
17+
18+ ############
19+ # Solution #
20+ ############
21+
22+ def find_k_closes (arr , pt , k ):
23+ n = len (arr )
24+ if k > n :
25+ return None
26+ if k < 1 :
27+ return None
28+
29+ kth_closest (arr , k - 1 , 0 , n - 1 , pt )
30+
31+ return arr [:k ]
32+
33+ def kth_closest (arr , n , left , right , pt ):
34+ pivot = pivoting (arr , left , right , pt )
35+
36+ if pivot > n :
37+ kth_closest (arr , n , left , pivot - 1 , pt )
38+ elif pivot < n :
39+ kth_closest (arr , n , pivot + 1 , right , pt )
40+
41+ def pivoting (arr , left , right , pt ):
42+ # O(N) pivoting
43+ # takes the last element as pivot
44+ pivot_dist = sqr_dist (pt , arr [right ])
45+ new_pivot = left
46+
47+ # iterate the whole array (without the last element)
48+ # and put all elements closer than the pivot (last element) in the first K spots
49+ # with the new_pivot we're "counting" how many closer elements are there
50+ for j in range (left , right ):
51+ if sqr_dist (pt , arr [j ]) < pivot_dist :
52+ swap (arr , new_pivot , j )
53+ new_pivot += 1
54+
55+ # swap the last (pivot) element with the new_pivot position
56+ swap (arr , new_pivot , right )
57+
58+ # return the new pivot
59+ return new_pivot
60+
61+ def swap (arr , i , j ):
62+ # swaps two elements in an array
63+ temp = arr [i ]
64+ arr [i ] = arr [j ]
65+ arr [j ] = temp
66+
67+ def sqr_dist (a , b ):
68+ # no need from the square root
69+ return (a [0 ] - b [0 ]) ** 2 + (a [1 ] - b [1 ]) ** 2
70+
71+
72+
73+ ###########
74+ # Testing #
75+ ###########
76+
77+ # Test 1
78+ # Correct result => [(2, 2), (2, 1.5), (2, 1)]
79+ print (find_k_closes ([(0 , 1 ), (3 , 3 ), (1 , 2 ), (2 , 1.5 ), (3 , - 1 ), (2 , 1 ), (4 , 3 ), (5 , 1 ), (- 1 , 2 ), (2 , 2 )], (2 , 2 ), 3 ))
80+
81+ # Test 2
82+ # Correct result => [(1, 2), (2, 1)]
83+ print (find_k_closes ([(0 , 1 ), (2 , 1 ), (3 , 3 ), (1 , 2 )], (2 , 2 ), 2 ))
0 commit comments