Welcome, guest | Sign In | My Account | Store | Cart

Wraps bisect.bisect() in an easy to use class that supports key-functions and straight-forward search methods.

Python, 311 lines
  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
from bisect import bisect_left, bisect_right

class SortedCollection(object):
    '''Sequence sorted by a key function.

    SortedCollection() is much easier to work with than using bisect() directly.
    It supports key functions like those use in sorted(), min(), and max().
    The result of the key function call is saved so that keys can be searched
    efficiently.

    Instead of returning an insertion-point which can be hard to interpret, the
    five find-methods return a specific item in the sequence. They can scan for
    exact matches, the last item less-than-or-equal to a key, or the first item
    greater-than-or-equal to a key.

    Once found, an item's ordinal position can be located with the index() method.
    New items can be added with the insert() and insert_right() methods.
    Old items can be deleted with the remove() method.

    The usual sequence methods are provided to support indexing, slicing,
    length lookup, clearing, copying, forward and reverse iteration, contains
    checking, item counts, item removal, and a nice looking repr.

    Finding and indexing are O(log n) operations while iteration and insertion
    are O(n).  The initial sort is O(n log n).

    The key function is stored in the 'key' attibute for easy introspection or
    so that you can assign a new key function (triggering an automatic re-sort).

    In short, the class was designed to handle all of the common use cases for
    bisect but with a simpler API and support for key functions.

    >>> from pprint import pprint
    >>> from operator import itemgetter

    >>> s = SortedCollection(key=itemgetter(2))
    >>> for record in [
    ...         ('roger', 'young', 30),
    ...         ('angela', 'jones', 28),
    ...         ('bill', 'smith', 22),
    ...         ('david', 'thomas', 32)]:
    ...     s.insert(record)

    >>> pprint(list(s))         # show records sorted by age
    [('bill', 'smith', 22),
     ('angela', 'jones', 28),
     ('roger', 'young', 30),
     ('david', 'thomas', 32)]

    >>> s.find_le(29)           # find oldest person aged 29 or younger
    ('angela', 'jones', 28)
    >>> s.find_lt(28)           # find oldest person under 28
    ('bill', 'smith', 22)
    >>> s.find_gt(28)           # find youngest person over 28
    ('roger', 'young', 30)

    >>> r = s.find_ge(32)       # find youngest person aged 32 or older
    >>> s.index(r)              # get the index of their record
    3
    >>> s[3]                    # fetch the record at that index
    ('david', 'thomas', 32)

    >>> s.key = itemgetter(0)   # now sort by first name
    >>> pprint(list(s))
    [('angela', 'jones', 28),
     ('bill', 'smith', 22),
     ('david', 'thomas', 32),
     ('roger', 'young', 30)]

    '''

    def __init__(self, iterable=(), key=None):
        self._given_key = key
        key = (lambda x: x) if key is None else key
        decorated = sorted((key(item), item) for item in iterable)
        self._keys = [k for k, item in decorated]
        self._items = [item for k, item in decorated]
        self._key = key

    def _getkey(self):
        return self._key

    def _setkey(self, key):
        if key is not self._key:
            self.__init__(self._items, key=key)

    def _delkey(self):
        self._setkey(None)

    key = property(_getkey, _setkey, _delkey, 'key function')

    def clear(self):
        self.__init__([], self._key)

    def copy(self):
        return self.__class__(self, self._key)

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

    def __getitem__(self, i):
        return self._items[i]

    def __iter__(self):
        return iter(self._items)

    def __reversed__(self):
        return reversed(self._items)

    def __repr__(self):
        return '%s(%r, key=%s)' % (
            self.__class__.__name__,
            self._items,
            getattr(self._given_key, '__name__', repr(self._given_key))
        )

    def __reduce__(self):
        return self.__class__, (self._items, self._given_key)

    def __contains__(self, item):
        k = self._key(item)
        i = bisect_left(self._keys, k)
        j = bisect_right(self._keys, k)
        return item in self._items[i:j]

    def index(self, item):
        'Find the position of an item.  Raise ValueError if not found.'
        k = self._key(item)
        i = bisect_left(self._keys, k)
        j = bisect_right(self._keys, k)
        return self._items[i:j].index(item) + i

    def count(self, item):
        'Return number of occurrences of item'
        k = self._key(item)
        i = bisect_left(self._keys, k)
        j = bisect_right(self._keys, k)
        return self._items[i:j].count(item)

    def insert(self, item):
        'Insert a new item.  If equal keys are found, add to the left'
        k = self._key(item)
        i = bisect_left(self._keys, k)
        self._keys.insert(i, k)
        self._items.insert(i, item)

    def insert_right(self, item):
        'Insert a new item.  If equal keys are found, add to the right'
        k = self._key(item)
        i = bisect_right(self._keys, k)
        self._keys.insert(i, k)
        self._items.insert(i, item)

    def remove(self, item):
        'Remove first occurence of item.  Raise ValueError if not found'
        i = self.index(item)
        del self._keys[i]
        del self._items[i]

    def find(self, k):
        'Return first item with a key == k.  Raise ValueError if not found.'
        i = bisect_left(self._keys, k)
        if i != len(self) and self._keys[i] == k:
            return self._items[i]
        raise ValueError('No item found with key equal to: %r' % (k,))

    def find_le(self, k):
        'Return last item with a key <= k.  Raise ValueError if not found.'
        i = bisect_right(self._keys, k)
        if i:
            return self._items[i-1]
        raise ValueError('No item found with key at or below: %r' % (k,))

    def find_lt(self, k):
        'Return last item with a key < k.  Raise ValueError if not found.'
        i = bisect_left(self._keys, k)
        if i:
            return self._items[i-1]
        raise ValueError('No item found with key below: %r' % (k,))

    def find_ge(self, k):
        'Return first item with a key >= equal to k.  Raise ValueError if not found'
        i = bisect_left(self._keys, k)
        if i != len(self):
            return self._items[i]
        raise ValueError('No item found with key at or above: %r' % (k,))

    def find_gt(self, k):
        'Return first item with a key > k.  Raise ValueError if not found'
        i = bisect_right(self._keys, k)
        if i != len(self):
            return self._items[i]
        raise ValueError('No item found with key above: %r' % (k,))


# ---------------------------  Simple demo and tests  -------------------------
if __name__ == '__main__':

    def ve2no(f, *args):
        'Convert ValueError result to -1'
        try:
            return f(*args)
        except ValueError:
            return -1

    def slow_index(seq, k):
        'Location of match or -1 if not found'
        for i, item in enumerate(seq):
            if item == k:
                return i
        return -1

    def slow_find(seq, k):
        'First item with a key equal to k. -1 if not found'
        for item in seq:
            if item == k:
                return item
        return -1

    def slow_find_le(seq, k):
        'Last item with a key less-than or equal to k.'
        for item in reversed(seq):
            if item <= k:
                return item
        return -1

    def slow_find_lt(seq, k):
        'Last item with a key less-than k.'
        for item in reversed(seq):
            if item < k:
                return item
        return -1

    def slow_find_ge(seq, k):
        'First item with a key-value greater-than or equal to k.'
        for item in seq:
            if item >= k:
                return item
        return -1

    def slow_find_gt(seq, k):
        'First item with a key-value greater-than or equal to k.'
        for item in seq:
            if item > k:
                return item
        return -1

    from random import choice
    pool = [1.5, 2, 2.0, 3, 3.0, 3.5, 4, 4.0, 4.5]
    for i in range(500):
        for n in range(6):
            s = [choice(pool) for i in range(n)]
            sc = SortedCollection(s)
            s.sort()
            for probe in pool:
                assert repr(ve2no(sc.index, probe)) == repr(slow_index(s, probe))
                assert repr(ve2no(sc.find, probe)) == repr(slow_find(s, probe))
                assert repr(ve2no(sc.find_le, probe)) == repr(slow_find_le(s, probe))
                assert repr(ve2no(sc.find_lt, probe)) == repr(slow_find_lt(s, probe))
                assert repr(ve2no(sc.find_ge, probe)) == repr(slow_find_ge(s, probe))
                assert repr(ve2no(sc.find_gt, probe)) == repr(slow_find_gt(s, probe))
            for i, item in enumerate(s):
                assert repr(item) == repr(sc[i])        # test __getitem__
                assert item in sc                       # test __contains__ and __iter__
                assert s.count(item) == sc.count(item)  # test count()
            assert len(sc) == n                         # test __len__
            assert list(map(repr, reversed(sc))) == list(map(repr, reversed(s)))    # test __reversed__
            assert list(sc.copy()) == list(sc)          # test copy()
            sc.clear()                                  # test clear()
            assert len(sc) == 0

    sd = SortedCollection('The quick Brown Fox jumped'.split(), key=str.lower)
    assert sd._keys == ['brown', 'fox', 'jumped', 'quick', 'the']
    assert sd._items == ['Brown', 'Fox', 'jumped', 'quick', 'The']
    assert sd._key == str.lower
    assert repr(sd) == "SortedCollection(['Brown', 'Fox', 'jumped', 'quick', 'The'], key=lower)"
    sd.key = str.upper
    assert sd._key == str.upper
    assert len(sd) == 5
    assert list(reversed(sd)) == ['The', 'quick', 'jumped', 'Fox', 'Brown']
    for item in sd:
        assert item in sd
    for i, item in enumerate(sd):
        assert item == sd[i]
    sd.insert('jUmPeD')
    sd.insert_right('QuIcK')
    assert sd._keys ==['BROWN', 'FOX', 'JUMPED', 'JUMPED', 'QUICK', 'QUICK', 'THE']
    assert sd._items == ['Brown', 'Fox', 'jUmPeD', 'jumped', 'quick', 'QuIcK', 'The']
    assert sd.find_le('JUMPED') == 'jumped', sd.find_le('JUMPED')
    assert sd.find_ge('JUMPED') == 'jUmPeD'
    assert sd.find_le('GOAT') == 'Fox'
    assert sd.find_ge('GOAT') == 'jUmPeD'
    assert sd.find('FOX') == 'Fox'
    assert sd[3] == 'jumped'
    assert sd[3:5] ==['jumped', 'quick']
    assert sd[-2] == 'QuIcK'
    assert sd[-4:-2] == ['jumped', 'quick']
    for i, item in enumerate(sd):
        assert sd.index(item) == i
    try:
        sd.index('xyzpdq')
    except ValueError:
        pass
    else:
        assert 0, 'Oops, failed to notify of missing value'
    sd.remove('jumped')
    assert list(sd) == ['Brown', 'Fox', 'jUmPeD', 'quick', 'QuIcK', 'The']

    import doctest
    from operator import itemgetter
    print(doctest.testmod())

For many uses of bisect(), this module is much easier to use.

The bisect module returns insertion-points which can be hard to interpret when searching. This module uses simple notions of finding exact matches, finding less-than-equal, or finding greater-than-equal.

This class also supports key-functions in a smart way (where the key function gets called no more than once per item).

Slicing and indexing are also supported.

10 comments

Mat Mathews 12 years, 1 month ago  # | flag

I think this is great!

I haven't grok'd it all exactly, maybe because 'key' as the itemgetter in the __init__ seems like confusing language. Why not have a itemgetter arg in the _le, _ge functions? keyIndex=

I might make a stab at adding namedtuples as the value in the collection with searchable field names, add itemgetter into the func signatures, add Set union and intersect functions for aggregated or joined results.

This SortedCollection could be a nice way to find, sort and slice in tuple space.

Tasty Recipe.

Joshua Smith 11 years, 8 months ago  # | flag

In the find_ge function, the if statement below (at line 167 above) seems incorrect to me. If the first key is .2, and I do find_ge(0.0), it should give me the item with key .2. Because of this if statement, it generates an error instead.

Am I missing something? Is there some way this could possibly be correct?

Offending code:

 if i == 0:
        raise ValueError('No item found with key at or above: %r' % (key,))
Raymond Hettinger (author) 11 years, 8 months ago  # | flag

Joshua, thanks for the comment. I've fixed that bug and added more tests.

Mat, the "key" specification is in __init__ because typical use cases set the key function just once and all of the calls to find, find_le, and find_ge use that same function.

Daniel Stutzbach 11 years, 8 months ago  # | flag

For a similar class with asymptotically faster inserts, see the sortedlist and sortedset classes in my blist package.

Alexander M Beedie 9 years, 8 months ago  # | flag

Slicing support can be improved; should really return another SortedCollection, not a list. Simplest way to do that extends __getitem__:

def __getitem__( self, i ):
    if isinstance( i,slice ):
        sc = self.__class__( key=self._key )
        sc._keys  = self._keys[i]
        sc._items = self._items[i]
        return sc
    else:
        return self._items[i]
Paul Hofstra 8 years, 6 months ago  # | flag

Raymond, I like this class a lot. I am missing 2 methods though: index_le and index_ge. index_le: return index of the largest item, that is less or equal to item and similar for index_ge. So probably also index_lt and index_gt might be useful.

Grant Jenks 8 years, 2 months ago  # | flag

If you want to see this idea taken further (pure-Python), made faster (fast-as-C implementations), fully tested (100% coverage and stress) and documented (with performance comparison) then check out the [sortedcontainers](http://grantjenks.com/docs/sortedcontainers/) available on [PyPI](https://pypi.python.org/pypi/sortedcontainers) and [github](https://github.com/grantjenks/sorted_containers).

Grant Jenks 8 years, 2 months ago  # | flag

If you want to see this idea taken further (pure-Python), made faster (fast-as-C implementations), fully tested (100% coverage and stress) and documented (with performance comparison) then check out the sortedcontainers available on PyPI and github.

Terry Jones 7 years, 2 months ago  # | flag

Hi Raymond

This is going to sound impossible (and I think it should be impossible), but...

I've been using this code recently (thanks!) and have run across an interesting issue. I am trying to make a small test case, but haven't managed yet. Anyway, I have a list of class instances (a class I wrote) each of which has numeric attribute. I pass the list and an attrgetter key to your class and I don't always get the same value 'decorated'.

To try to prove this to myself, I print both the original iterable and also the decorated one. In both calls, the original is the same but the decorated version is different. To be doubly sure I also calculate md5 sums on str versions of iterable and on the decorated - same result.

I know Python sort is supposed to be stable. Like everyone else, I would even say it IS stable. But if so, I can't see how I'm getting this behavior. In one call, I get a decorated variable with these 3 contiguous elements:

 (38, Prosite symbol='PS' offset=38 len=4 detail='00294')
 (38, Prosite symbol='PS' offset=38 len=3 detail='00342')
 (38, AminoAcidsLm symbol='N' offset=38 len=1 detail='')

and in another call, with the same input, decorated contains those 3 elements in a different order. Still contiguous. So the sort is working fine on the offset (key), but the rest of it seems non-stable (BTW, those strings like "Prosite symbol='PS' offset=38 len=4 detail='00294'" are just coming from my class' __str__ method). My class doesn't define __lt__ or other comparison functions (though it does have an __eq__), but I don't think that should make any difference, given Python's sort stability.

This is all happening within one Python (2.7.6) process.

But.... when I take the input data I'm passing to your class and I simply make my own 'decorated' variable from it, copying your __init__ code, the decorated result is always identical.

I.e., I haven't managed to make this fail in a simplified standalone example. Frustrating!

Have you heard of anything like this? Or do you have any thoughts? I've been in the game long enough to know that there's a 99.99% chance I'm in the wrong somehow, but I can't see how. Surely my "proof" that the class __init__ is being called with the same thing in both cases would lead to a reasonable expectation that I'd get the same result. But I don't.

I'll keep trying to get a simple failing case. That's probably the best way to find out what I'm doing wrong / don't understand. It's weird, I haven't debugged like this in a loooong time! :-)

Thanks for any help!

Regards, Terry

Raymond Hettinger (author) 7 years, 2 months ago  # | flag

Hello Terry,

I think your sort stability has been confounded by the lexicographic sort order which is ordering the tuples by the decorated value as the primary key and the original value as the secondary key.

To preserve the stability of the input, make the following modification to the recipe:

    decorated = sorted((key(item), i, item) for i, item in enumerate(iterable))
    self._keys = [k for k, i, item in decorated]
    self._items = [item for k, i, item in decorated]