heapq avec prédicat de comparaison personnalisé

Question

J'essaie de construire un tas avec un prédicat de tri personnalisé. Comme les valeurs qui y entrent sont de type "défini par l'utilisateur", je ne peux pas modifier leur prédicat de comparaison intégré.

Y a-t-il un moyen de faire quelque chose comme :

h = heapq.heapify([...], key=my_lt_pred)
h = heapq.heappush(h, key=my_lt_pred)

Ou encore mieux, je pourrais envelopper les fonctions heapq dans mon propre conteneur afin de ne pas avoir à passer le prédicat.

Answer 1

En python3, vous pouvez utiliser cmp_to_key de functools module. code source cpython .

Supposons que vous ayez besoin d'une file d'attente prioritaire de triplets et que vous spécifiez la priorité en utilisant le dernier attribut.

from heapq import *
from functools import cmp_to_key
def mycmp(triplet_left, triplet_right):
    key_l, key_r = triplet_left[2], triplet_right[2]
    if key_l > key_r:
        return -1  # larger first
    elif key_l == key_r:
        return 0  # equal
    else:
        return 1

WrapperCls = cmp_to_key(mycmp)
pq = []
myobj = tuple(1, 2, "anystring")
# to push an object myobj into pq
heappush(pq, WrapperCls(myobj))
# to get the heap top use the `obj` attribute
inner = pq[0].obj

Test de performance :

Environnement

python 3.10.2

Code

from functools import cmp_to_key
from timeit import default_timer as time
from random import randint
from heapq import *

class WrapperCls1:
    __slots__ = 'obj'
    def __init__(self, obj):
        self.obj = obj
    def __lt__(self, other):
        kl, kr = self.obj[2], other.obj[2]
        return True if kl > kr else False

def cmp_class2(obj1, obj2):
    kl, kr = obj1[2], obj2[2]
    return -1 if kl > kr else 0 if kl == kr else 1

WrapperCls2 = cmp_to_key(cmp_class2)

triplets = [[randint(-1000000, 1000000) for _ in range(3)] for _ in range(100000)]
# tuple_triplets = [tuple(randint(-1000000, 1000000) for _ in range(3)) for _ in range(100000)]

def test_cls1():
    pq = []
    for triplet in triplets:
        heappush(pq, WrapperCls1(triplet))

def test_cls2():
    pq = []
    for triplet in triplets:
        heappush(pq, WrapperCls2(triplet))

def test_cls3():
    pq = []
    for triplet in triplets:
        heappush(pq, (-triplet[2], triplet))

start = time()
for _ in range(10):
    test_cls1()
    # test_cls2()
    # test_cls3()
print("total running time (seconds): ", -start+(start:=time()))

Résultats

utiliser list au lieu de tuple par fonction :

• WrapperCls1 : 16.2ms
• WrapperCls1 avec __slots__ : 9.8ms
• WrapperCls2 : 8.6ms
• déplacer l'attribut de priorité en première position ( ne pas supporter personnalisé prédicat ) : 6.0ms.

Par conséquent, cette méthode est un peu plus rapide que l'utilisation d'une classe personnalisée avec un paramètre surchargé __lt__() et la fonction __slots__ attribut.