Thursday, September 23, 2010

[Python]几个常用数据结构的高效实现

Python中最广为熟知的container就是list和dict了. 但是对于特定的用途, 这二者未必是最高效的.比如判断容器中是否包括x(x in container),如果容器是list就需要O(n)的时间.而使用set的话这个就只需要O(1)的时间. 我的血泪教训是: 在list很长的时候, x in list这个操作是非常非常低效的.一个好的主意是在另一个set容器来保存membership的信息.

集合(Set)
set是Python的built-in类型,对于membership testing, removing duplicates 的操作非常高校(O(1)时间).另外对集合操作也非常方便.如果对容器元素的顺序没有要求,set是非常好的选择对象
>>>s= set()
>>>s.add(1)
>>>s.add(2)
>>>s.remove(1)  #throw exception if not in set
>>>s.discard(1) #remove if present
>>>print set([1,2,3]) | set([3,4,5]) #union
set([1, 2, 3, 4, 5])
>>>print set([1,2,3]) & set([3,4,5]) #intersection
set([1])
>>>print set([1,2,3]) - set([3,4,5]) #difference
set([1, 2])
>>>print set([1,2,3]) ^ set([3,4,5]) #symetric difference
set([1, 2, 4, 5])
栈(Stack)
栈的特点是先进后出(FILO).built-in的list用来实现栈非常合适.对于append和pop的实现非常高效.

FIFO队列(FIFO Queue)
队列的特点是先进先出(FIFO). list用来实现队列不够有效.collections.deque可以弥补这一缺点(fast appends and pops from both ends)
>>>from collections import deque
>>>queue = deque()
>>>queue.append(1)
>>>queue.append(2)
>>>queue.appendleft(3)
>>>print queue
deque([3, 1, 2])
>>>queue.pop()
2
>>>queue.popleft()
3

利用deque,我们可以实现一个时间上高效的FIFO queue:
class FIFOQueue():
    def __init__(self):
        self.queue = deque()
        self.counter = itertools.count(1)    # unique sequence count
        self.task_finder = {}                # mapping of tasks to entries

    def append(self, task):
        if debug:
            print "FIFOQueue.append "+repr(task)

        if task in self:
            return
        count = next(self.counter)

        entry = [count, task]
        self.task_finder[task] = entry
        self.queue.append(entry)
        if debug:
            print "FIFOQueue.append "+repr(task)+" done"

    def top(self):
        if debug:
            print "FIFOQueue.top "
        while self.queue:
            count, task  = self.queue[0]
            if count:
                if debug:
                    print "FIFOQueue.pop "+repr(task)+" done"
                return task
            else:
                self.queue.popleft()
        if debug:
            print "FIFOQueue.top none"
        return None

    def pop(self):
        if debug:
            print "FIFOQueue.pop "
        while not self.top():
            continue
        if self.queue:
            entry  = self.queue.popleft()
            count, task = entry
            assert(count != None)
            assert(self.task_finder[task] == entry)
            del self.task_finder[task]
            if debug:
                print "FIFOQueue.pop "+repr(task)+" done"
            return task
        if debug:
            print "FIFOQueue.pop none"
        return None

    def remove(self, task):
        if debug:
            print "FIFQueue.remove "+repr(task)
        if task not in self:
            return
        entry = self.task_finder[task]
        assert(entry[0] != None)
        entry[0] = None
        del self.task_finder[task]
        if debug:
            print "FIFOQueue.remove "+repr(task)+" done"

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

    def __contains__(self, task):
        return self.task_finder.has_key(task)

    def __repr__(self):
        return repr(self.queue)
这个实现的好处在于可以常数时间删除一个元素以及在常数时间内判定一个元素是否属于这个queue

优先队列(Priority Queue)
使用堆(heap)实现的优先队列(priority queue)可以在O(1)时间内得到队列中最小元素,O(logn)时间内插入或者删除一个元素.我通常在simulator中会使用priority queue来管理所有事件.

Python提供了heapq来实现.值得注意的是heapq并不是一个容器,而是操作在list上的一系列的heap算法(包括heappush,heappop,heapify)等.http://docs.python.org/library/heapq.html给出了用这一些算法实现堆排序(heapsort)以及优先队列的. 不过在那个优先队列的例子里有一点小bug. 这里是我给出的代码
import itertools
from heapq import heappush, heappop
class PriorityQueue():
    def __init__(self):
        self.pq = []                         # the priority queue list
        self.counter = itertools.count(1)    # unique sequence count
        self.task_finder = {}                # mapping of tasks to entries
        self.size = 0
 
    def add_task(self, priority, task, count=None):
        if count is None:
            count = next(self.counter)
        entry = [priority, count, task]
        self.task_finder[task] = entry
        heappush(self.pq, entry)
        self.size  += 1

    def pop_task(self):
        while self.pq:
            entry = heappop(self.pq)
            priority, count, task = entry
            if self.task_finder[task] == entry:
                del self.task_finder[task]
            if count is not None:
                self.size -= 1
                return (priority, task)
        return None


    def del_task(self, task):
        entry = self.task_finder[task]
        if entry[1] != None:
            self.size -= 1
        entry[1] = None

    def reprioritize(self, priority, task):
        entry = self.task_finder[task]
        self.add_task(priority, task, entry[1])
        if entry[1] != None:
            self.size -= 1
        entry[1] = None

    def __len__(self):
        return self.size

No comments: