import heapq
import random


class Solution(object):
    def kClosest(self, points, K):
        """
        :type points: List[List[int]]
        :type K: int
        :rtype: List[List[int]]
        """
        arr = []

        def distance(p):
            return p[0] ** 2 + p[1] ** 2

        for p in points:
            t = distance(p)
            if len(arr) < K:
                heapq.heappush(arr, (-t, p))
            elif -t > arr[0][0]:
                heapq.heappop(arr)
                heapq.heappush(arr, (-t, p))

        res = [p for t, p in arr]

        return res

    def kClosest_partition(self, points, K):

        """
        :type points: List[List[int]]
        :type K: int
        :rtype: List[List[int]]
        """
        dist = lambda i: points[i][0] ** 2 + points[i][1] ** 2

        def partition(left, right, n):
            if left >= right:
                return
            idx = random.randint(left, right)
            points[idx], points[right] = points[right], points[idx]
            pivot, k = dist(idx), left

            for t in xrange(left, right + 1):
                if dist(t) < pivot:
                    points[k], points[t] = points[t], points[k]
                    k += 1
            points[right], points[k] = points[k], points[right]

            if k <= n:
                return partition(k, right, n - k)
            else:
                return partition(left, k - 1, n)

        partition(0, len(points) - 1, K)
        return points[:K]

    def kClosest2(self, points, K):
        """
        :type points: List[List[int]]
        :type K: int
        :rtype: List[List[int]]
        """

        dist = lambda i: points[i][0] ** 2 + points[i][1] ** 2

        def work(i, j, K):
            if i >= j: return
            oi, oj = i, j
            pivot = dist(random.randint(i, j))
            while i < j:
                while i < j and dist(i) < pivot: i += 1
                while i < j and dist(j) > pivot: j -= 1
                points[i], points[j] = points[j], points[i]

            if K <= i - oi + 1:
                work(oi, i, K)
            else:
                work(i + 1, oj, K - (i - oi + 1))

        work(0, len(points) - 1, K)
        return points[:K]





if __name__ == '__main__':
    print Solution().kClosest([[3, 3], [5, -1], [-2, 4]], 2)
    print Solution().kClosest_partition([[3, 3], [5, -1], [-2, 4]], 2)
    print Solution().kClosest([[10, -2], [2, -2], [10, 10], [9, 4], [-8, 1]], 4)
    print Solution().kClosest_partition([[10, -2], [2, -2], [10, 10], [9, 4], [-8, 1]], 4)
    print Solution().kClosest([[1, 3], [-2, 2]], 1)
    print Solution().kClosest_partition([[1, 3], [-2, 2]], 1)
    print Solution().kClosest([[3, 3], [5, -1], [-2, 4]], 2)
    print Solution().kClosest_partition([[3, 3], [5, -1], [-2, 4]], 2)
    print Solution().kClosest(
        [[-95, 76], [17, 7], [-55, -58], [53, 20], [-69, -8], [-57, 87], [-2, -42], [-10, -87], [-36, -57], [97, -39],
         [97, 49]], 5)
    print Solution().kClosest_partition(
        [[-95, 76], [17, 7], [-55, -58], [53, 20], [-69, -8], [-57, 87], [-2, -42], [-10, -87], [-36, -57], [97, -39],
         [97, 49]], 5)
