K Smallest Elements - Heap v/s Sort

import perfplot
import heapq
import numpy as np

key = lambda P: P[0]**2 + P[1]**2

def sort(inp):
    l, n = inp
    l = l.copy()
    l.sort(key=key)
    return l[:n]

def heap(inp):
    l, n = inp
    return heapq.nsmallest(n, l, key=key)

perfplot.show(
    setup=lambda n: (np.random.randint(-n, n, (n, 2)).tolist(), int(n * .2)),
    kernels=[heap, sort],
    labels=['heapq', 'timsort'],
    n_range=[2 ** k for k in range(4, 22)],
    xlabel='N',
)