Python heapq:如何使用列表列表的第 n 个元素对堆进行排序?

Adi*_*rel 4 python

所以我有列表被添加到堆中;例如:

 n = [[1, 5, 93],
     [2, 6, 44],
     [4, 7, 45],
     [6, 3, 12]]

 heapq.heapify(n)
 print(n)
Run Code Online (Sandbox Code Playgroud)

这将根据列表的第一个元素进行比较和排序。

我的问题是,如何对 heapq 进行排序,以便比较每个列表的第三个元素?例如,上面的列表将按以下顺序从 heapq 访问:

[[6, 3, 12],
 [2, 6, 44],
 [4, 7, 45],
 [1, 5, 93]]
Run Code Online (Sandbox Code Playgroud)

ACh*_*ion 8

heapq不支持key它的排序函数,因此您需要操作您的数据结构。tuple(sort_value, list)将您的列表映射到 a将允许您进行log(n)推送和弹出:

 In []:
 q = [(x[2], x) for x in n]
 heapq.heapify(q)
 heapq.heappop(q)

 Out[]:
 (12, [6, 3, 12])

 In []:
 l = [2, 5, 1]
 heapq.heappush(q, (l[2], l))
 heapq.heappop(q)

 Out[]:
 (1, [2, 5, 1])
Run Code Online (Sandbox Code Playgroud)

或者,定义您自己的list并为该列表实现比较函数:

class MyList(list):
    def __lt__(self, other):
        return self[2] < other[2]

q = [MyList(x) for x in n]
Run Code Online (Sandbox Code Playgroud)

注意:您应该实现其他比较函数(参见functools.total_ordering如何轻松实现)。