python实现堆和索引堆的代码示例

堆是一棵完全二叉树。堆分为大根堆和小根堆，大根堆是父节点大于左右子节点，并且左右子树也满足该性质的完全二叉树。小根堆相反。可以利用堆来实现优先队列。

由于是完全二叉树，所以可以使用数组来表示堆，索引从0开始[0:length-1]。结点i的左右子节点分别为2i+1,2i+2。长度为length的树的最后一个非叶子节点为length//2-1。当前节点i的父节点为(i-1)//2。其中//表示向下取整。

以大根堆举例。当每次插入或者删除的时候，为了保证堆的结构特征不被破坏，需要进行调整。调整分为两种，一种是从上往下，将小的数下沉。一种是从下往上，令大的数上浮。

具体实现如下：

首先编写几个魔术方法。包括构造函数，可以直接调用len来返回data数组长度的函数，一个打印data内容的函数

def __init__(self, data=[]):
   self.data = data
   self.construct_heap()

def __len__(self):
   return len(self.data)

def __str__(self):
   return str(self.data)

定义一个swap函数，来方便的交换数组中两个索引处的值。

 def swap(self, i, j):
   self.data[i], self.data[j] = self.data[j], self.data[i]

定义float_up方法，使堆中大的数能浮上来。当前节点不为根节点，并且当前节点数据大小大于父节点时，上浮。

def float_up(self, i):
   while i > 0 and self.data[i] > self.data[(i - 1) // 2]:
     self.swap(i, (i - 1) // 2)
     i = (i - 1) // 2

定义sink_down方法，使堆中小的数沉下去。当前节点不为叶子节点时，如果小于左孩子或右孩子的数据，则和左右孩子中较大的换一下位置。

def sink_down(self, i):
   while i < len(self) // 2:
     l, r = 2 * i + 1, 2 * i + 2
     if r < len(self) and self.data[l] < self.data[r]:
       l = r
     if self.data[i] < self.data[l]:
       self.swap(i, l)

i = l

实现append方法，能够动态地添加数据。在数据数组尾部添加数据，然后将数据上浮。

 def append(self, data):
   self.data.append(data)
   self.float_up(len(self) - 1)

实现pop_left方法，取堆中最大元素，即优先队列中第一个元素。将数组中第一个元素与最后一个元素换位置，删除最后一个元素，然后将第一个元素下沉到合适的位置。

 def pop_left(self):
   self.swap(0, len(self) - 1)
   r = self.data.pop()
   self.sink_down(0)
   return r

如果想在初始化堆的时候，向构造函数中传入数据参数，则需要一次性将整个堆构建完毕，而不能一个一个加入。实现也很简单，从最后一个非叶节点开始，逐个执行sink_down操作。

 def construct_heap(self):
   for i in range(len(self) // 2 - 1, -1, -1):
     self.sink_down(i)

这样一个基本的堆的代码就编写完毕了。

但是如果我们想要动态的改变数据，当前的堆就不能满足我们的需求了，因为索引不能总是标识同一个数据，因为堆的结构是不断调整的。我们需要使用索引堆。

在索引堆中，我们不在堆中直接保存数据，而是用在堆中存放数据的索引。

如果我们输入的数据arr是 45 20 12 5 35。则arr[0]一直指向45，arr[1]一直指向20，因为我们在调整堆结构中实际调整的是索引数组，而不会改变真实存放数据的数组。

因此我们的代码需要调整，首先在构造函数中加入一个索引数组。下标从0开始，与存放数据的数组的下标相对应。

 def __init__(self, data=[]):
   self.data = data
   self.index_arr = list(range(len(self.data)))
   self.construct_heap()

然后将返回堆长度的魔术函数也修改一下。

 def __len__(self):
   return len(self.index_arr)

调整一下之前定义的swap方法，原来是直接交换数据，现在交换索引。

def swap(self, i, j):
   self.index_arr[i], self.index_arr[j] = self.index_arr[j], self.index_arr[i]

调整float_up以及sink_down中的相应位置

 def float_up(self, i):
   while i > 0 and self.data[self.index_arr[i]] > self.data[self.index_arr[(i - 1) // 2]]:
     self.swap(i, (i - 1) // 2)
     i = (i - 1) // 2

def sink_down(self, i):
   while i < len(self) // 2:
     l, r = 2 * i + 1, 2 * i + 2
     if r < len(self) and self.data[self.index_arr[l]] < self.data[self.index_arr[r]]:
       l = r
     if self.data[self.index_arr[i]] < self.data[self.index_arr[l]]:
       self.swap(i, l)

i = l

当append数据的时候，要相应的更新index_arr

 def append(self, data):
   self.data.append(data)
   self.index_arr.append(len(self))
   self.float_up(len(self) - 1)

当移出数据的时候，之前已经提到过存放数据的数组，是按照append的顺序进行存储的，平时操作只是对index_arr的顺序进行调整。

如果data_arr为 42 30 74 60 相应的index_arr应该为2 3 0 1

这时，当我们popleft出最大元素时，data_arr中的74被移出后变成了42 30 60，数组中最大索引由3变成了2，如果索引数组中仍然用3这个索引来索引30会造成index溢出。74的索引为2，需要我们将索引数在2之后的都减1。

综上，在删除元素时，我们原先是将data_arr中的首尾元素互换，再删除尾部元素，再对头部元素进行sink_down操作。现在我们先换索引数组中首尾元素，再删除索引数组尾部元素，此时尚未操作存放data的data_arr，因此索引数组剩余元素与data_arr的元素仍是一一对应的。进行sink_down操作，操作完成之后再删除data_arr相应位置元素。最后将index_arr中值大于原index_arr头部元素值的减一。

 def pop_left(self):
   self.swap(0, len(self) - 1)
   r = self.index_arr.pop()
   self.sink_down(0)
   self.data.pop(r)

for i, index in enumerate(self.index_arr):
     if index > r:
       self.index_arr[i] -= 1

return r

索引堆增加了一个更新操作，可以随时更新索引堆中的数据。更新时，先直接更新data_arr中相应索引处的数据，然后在index_arr中，找到存放了data_arr中，刚被更新的数据的索引的索引位置，与删除时一样需要进行一次遍历。找到这个位置之后，由于无法确定与前后元素的大小关系，因此需要进行一次float_up操作再进行一次sink_down操作。

 def update(self, i, data):
   self.data[i] = data
   for index_index, index in enumerate(self.index_arr):
     if index == i:
       target = index_index

self.float_up(target)
   self.sink_down(target)

可以很明显看出，这个索引堆在插入元素时是比较快的，但是在删除元素和更新元素时，为了查找相应位置索引，都进行了一次遍历，这是很耗时的操作。为了能更快的找到index_arr中值为要更新的data_arr的相应索引值得索引位置，我们再次开辟一个新的数组to_index，来对index_arr进行索引。

例如对于数组75 54 65 90

此时它的index_arr为3 0 2 1。当要更新data[3]，即90这个元素时，现在要遍历一遍index_arr来找到3这个位置，这个位置是0。我们要建立一个to_index，to_index[3]中存放的元素为0。

index_arr存放的元素分别为： 1 3 2 0。

先改变swap数组，在交换index_arr中元素时，也交换存放在to_index中的index_arr的索引。

def swap(self, i, j):
   self.index_arr[i], self.index_arr[j] = self.index_arr[j], self.index_arr[i]
   self.to_index[self.index_arr[i]], self.to_index[self.index_arr[j]] = self.to_index[self.index_arr[j]], \
                                      self.to_index[self.index_arr[i]]

然后在update中，当要更新位置为i的元素时，我们就不需要通过一次遍历才能找到index_arr中该元素的索引，而是直接通过访问index_arr[i]即可访问index_arr中相应索引

 def update(self, i, data):
   self.data[i] = data
   target = self.to_index[i]
   self.float_up(target)
   self.sink_down(target)

最后改变pop_left中相应代码，这时我们需要维护三个数组，data_arr，index_arr以及to_index。

仍然是首先将index_arr首位元素交换，并pop出尾部元素存放到i中。然后将头部元素sink_down到相应位置，然后将pop出data_arr索引i处的元素。然后pop出to_index中索引为i的元素，再将index_arr中索引溢出的元素进行调整。

 def pop_left(self):
   self.swap(0, len(self) - 1)
   r = self.index_arr.pop()

self.sink_down(0)
   self.data.pop(r)

self.to_index.pop(r)
   for i in range(r, len(self)):
     self.index_arr[self.to_index[i]] -= 1

return r

来源：https://segmentfault.com/a/1190000013774311