Forma rápida de obtener elementos N Mín. O Máx. De una lista en Python

Question

Actualmente tengo una larga lista que se ordena usando una función lambda f. Luego elijo un elemento aleatorio de los primeros cinco elementos. Algo así como:

f = lambda x: some_function_of(x, local_variable) 
my_list.sort(key=f) 
foo = choice(my_list[:4]) 

Esto es un cuello de botella en mi programa, de acuerdo con el generador de perfiles. ¿Cómo puedo acelerar las cosas? ¿Existe una manera rápida e incorporada de recuperar los elementos que quiero (en teoría no debería ser necesario ordenar toda la lista)? Gracias.

Answer 1

Utilice heapq.nlargest o heapq.nsmallest.

Por ejemplo:

import heapq 

elements = heapq.nsmallest(4, my_list, key=f) 
foo = choice(elements) 

Esto se llevará a O (N + KlogN) tiempo (donde K es el número de elementos devueltos, y N es el tamaño de la lista), que es más rápida que O (NlogN) para el ordenamiento normal cuando K es pequeño en relación con N.

Answer 2

En realidad es posible en tiempo lineal (O (N)) en promedio.

Se necesita un algoritmo de partición:

def partition(seq, pred, start=0, end=-1): 
    if end == -1: end = len(seq) 
    while True: 
     while True: 
      if start == end: return start 
      if not pred(seq[start]): break 
      start += 1 
     while True: 
      if pred(seq[end-1]): break 
      end -= 1 
      if start == end: return start 
     seq[start], seq[end-1] = seq[end-1], seq[start] 
     start += 1 
     end -= 1 

que puede ser utilizado por un algoritmo nth_element:

def nth_element(seq_in, n, key=lambda x:x): 
    start, end = 0, len(seq_in) 
    seq = [(x, key(x)) for x in seq_in] 

    def partition_pred(x): return x[1] < seq[end-1][1] 

    while start != end: 
     pivot = (end + start) // 2 
     seq[pivot], seq[end - 1] = seq[end - 1], seq[pivot] 
     pivot = partition(seq, partition_pred, start, end) 
     seq[pivot], seq[end - 1] = seq[end - 1], seq[pivot] 
     if pivot == n: break 
     if pivot < n: start = pivot + 1 
     else: end = pivot 

    seq_in[:] = (x for x, k in seq) 

Teniendo en cuenta estos, simplemente reemplazar su segunda línea (más o menos) con:

nth_element(my_list, 4, key=f)