Statistics
| Branch: | Revision:

iof-tools / networkxMiCe / networkx-master / networkx / utils / mapped_queue.py @ 5cef0f13

History | View | Annotate | Download (6.15 KB)

1
# -*- coding: utf-8 -*-
2
#
3
# priorityq: An object-oriented priority queue with updatable priorities.
4
#
5
# Copyright 2018 Edward L. Platt
6
#
7
# This file is part of NetworkX
8
#
9
# NetworkX is distributed under a BSD license; see LICENSE.txt for more
10
# information.
11
#
12
# Authors:
13
#   Edward L. Platt <ed@elplatt.com>
14
#
15
"""Priority queue class with updatable priorities.
16
"""
17

    
18
import heapq
19

    
20
__all__ = ['MappedQueue']
21

    
22

    
23
class MappedQueue(object):
24
    """The MappedQueue class implements an efficient minimum heap. The
25
    smallest element can be popped in O(1) time, new elements can be pushed
26
    in O(log n) time, and any element can be removed or updated in O(log n)
27
    time. The queue cannot contain duplicate elements and an attempt to push an
28
    element already in the queue will have no effect.
29

30
    MappedQueue complements the heapq package from the python standard
31
    library. While MappedQueue is designed for maximum compatibility with
32
    heapq, it has slightly different functionality.
33

34
    Examples
35
    --------
36

37
    A `MappedQueue` can be created empty or optionally given an array of
38
    initial elements. Calling `push()` will add an element and calling `pop()`
39
    will remove and return the smallest element.
40

41
    >>> q = MappedQueue([916, 50, 4609, 493, 237])
42
    >>> q.push(1310)
43
    True
44
    >>> x = [q.pop() for i in range(len(q.h))]
45
    >>> x
46
    [50, 237, 493, 916, 1310, 4609]
47

48
    Elements can also be updated or removed from anywhere in the queue.
49

50
    >>> q = MappedQueue([916, 50, 4609, 493, 237])
51
    >>> q.remove(493)
52
    >>> q.update(237, 1117)
53
    >>> x = [q.pop() for i in range(len(q.h))]
54
    >>> x
55
    [50, 916, 1117, 4609]
56

57
    References
58
    ----------
59
    .. [1] Cormen, T. H., Leiserson, C. E., Rivest, R. L., & Stein, C. (2001).
60
       Introduction to algorithms second edition.
61
    .. [2] Knuth, D. E. (1997). The art of computer programming (Vol. 3).
62
       Pearson Education.
63
    """
64

    
65
    def __init__(self, data=[]):
66
        """Priority queue class with updatable priorities.
67
        """
68
        self.h = list(data)
69
        self.d = dict()
70
        self._heapify()
71

    
72
    def __len__(self):
73
        return len(self.h)
74

    
75
    def _heapify(self):
76
        """Restore heap invariant and recalculate map."""
77
        heapq.heapify(self.h)
78
        self.d = dict([(elt, pos) for pos, elt in enumerate(self.h)])
79
        if len(self.h) != len(self.d):
80
            raise AssertionError("Heap contains duplicate elements")
81

    
82
    def push(self, elt):
83
        """Add an element to the queue."""
84
        # If element is already in queue, do nothing
85
        if elt in self.d:
86
            return False
87
        # Add element to heap and dict
88
        pos = len(self.h)
89
        self.h.append(elt)
90
        self.d[elt] = pos
91
        # Restore invariant by sifting down
92
        self._siftdown(pos)
93
        return True
94

    
95
    def pop(self):
96
        """Remove and return the smallest element in the queue."""
97
        # Remove smallest element
98
        elt = self.h[0]
99
        del self.d[elt]
100
        # If elt is last item, remove and return
101
        if len(self.h) == 1:
102
            self.h.pop()
103
            return elt
104
        # Replace root with last element
105
        last = self.h.pop()
106
        self.h[0] = last
107
        self.d[last] = 0
108
        # Restore invariant by sifting up, then down
109
        pos = self._siftup(0)
110
        self._siftdown(pos)
111
        # Return smallest element
112
        return elt
113

    
114
    def update(self, elt, new):
115
        """Replace an element in the queue with a new one."""
116
        # Replace
117
        pos = self.d[elt]
118
        self.h[pos] = new
119
        del self.d[elt]
120
        self.d[new] = pos
121
        # Restore invariant by sifting up, then down
122
        pos = self._siftup(pos)
123
        self._siftdown(pos)
124

    
125
    def remove(self, elt):
126
        """Remove an element from the queue."""
127
        # Find and remove element
128
        try:
129
            pos = self.d[elt]
130
            del self.d[elt]
131
        except KeyError:
132
            # Not in queue
133
            raise
134
        # If elt is last item, remove and return
135
        if pos == len(self.h) - 1:
136
            self.h.pop()
137
            return
138
        # Replace elt with last element
139
        last = self.h.pop()
140
        self.h[pos] = last
141
        self.d[last] = pos
142
        # Restore invariant by sifting up, then down
143
        pos = self._siftup(pos)
144
        self._siftdown(pos)
145

    
146
    def _siftup(self, pos):
147
        """Move element at pos down to a leaf by repeatedly moving the smaller
148
        child up."""
149
        h, d = self.h, self.d
150
        elt = h[pos]
151
        # Continue until element is in a leaf
152
        end_pos = len(h)
153
        left_pos = (pos << 1) + 1
154
        while left_pos < end_pos:
155
            # Left child is guaranteed to exist by loop predicate
156
            left = h[left_pos]
157
            try:
158
                right_pos = left_pos + 1
159
                right = h[right_pos]
160
                # Out-of-place, swap with left unless right is smaller
161
                if right < left:
162
                    h[pos], h[right_pos] = right, elt
163
                    pos, right_pos = right_pos, pos
164
                    d[elt], d[right] = pos, right_pos
165
                else:
166
                    h[pos], h[left_pos] = left, elt
167
                    pos, left_pos = left_pos, pos
168
                    d[elt], d[left] = pos, left_pos
169
            except IndexError:
170
                # Left leaf is the end of the heap, swap
171
                h[pos], h[left_pos] = left, elt
172
                pos, left_pos = left_pos, pos
173
                d[elt], d[left] = pos, left_pos
174
            # Update left_pos
175
            left_pos = (pos << 1) + 1
176
        return pos
177

    
178
    def _siftdown(self, pos):
179
        """Restore invariant by repeatedly replacing out-of-place element with
180
        its parent."""
181
        h, d = self.h, self.d
182
        elt = h[pos]
183
        # Continue until element is at root
184
        while pos > 0:
185
            parent_pos = (pos - 1) >> 1
186
            parent = h[parent_pos]
187
            if parent > elt:
188
                # Swap out-of-place element with parent
189
                h[parent_pos], h[pos] = elt, parent
190
                parent_pos, pos = pos, parent_pos
191
                d[elt] = pos
192
                d[parent] = parent_pos
193
            else:
194
                # Invariant is satisfied
195
                break
196
        return pos