python3 数据结构和算法(2) 优先队列

import random
import time
from functools import wraps

def timethis(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        start = time.time()
        result = func(*args, **kwargs)
        end = time.time()
        print('%-30s %10s %f' % (func.__name__, ' cost time:', end - start))
        return result
    return wrapper


class PriorityQueue():
    def __init__(self, data):
        self.data = list(data)
        self.build_heap()

    def heapify(self, i, heap_size):
        l = (i+1) * 2 - 1
        r = (i+1) * 2

        largest = i
        if l < heap_size and self.data[l] > self.data[largest]:
            largest = l
        if r < heap_size and self.data[r] > self.data[largest]:
            largest = r

        if largest != i:
            self.data[i], self.data[largest] = self.data[largest], self.data[i]
            self.heapify(largest, heap_size)

    def build_heap(self):
        heap_size = len(self.data)
        for i in range(len(self.data)//2, -1, -1):
            self.heapify(i, heap_size)

    def extract(self):
        try:
            ret = self.data[0]
            self.data[0] = self.data[-1]
            self.data.pop()
            self.heapify(0, len(self.data))
            return ret
        except Exception:
            print('no more data')

    def insert(self, k):
        self.data.append(k)
        i = len(self.data) - 1
        while i > 0 and self.data[(i+1) // 2 - 1] < self.data[i]:
            self.data[i], self.data[(i+1) // 2 - 1] = self.data[(i+1) // 2 - 1], self.data[i]
            i = (i+1) // 2 - 1 

    def empty(self):
        return len(self.data) == 0

    def top(self):
        return self.data[0]


def test_insert():
    data = []
    pq = PriorityQueue(data)
    for i in range(100):
        pq.insert(random.randint(1, 1000))
        if pq.top() != max(pq.data):
            print('error')

    print()
    for i in range(100):
        print(pq.extract(), end=', ')
    print()

def main():
    data = [random.randint(1, 1000) for i in range(100)]
    pq = PriorityQueue(data)
    for i in range(50):
        print(pq.extract(), end=', ')
    print()

    test_insert()

if __name__ == '__main__':
    main()

猜你喜欢

转载自blog.csdn.net/guzhou_diaoke/article/details/81415452