SortedCollection(Python recipe)by Raymond Hettinger ActiveState Code (http://code.activestate.com/recipes/577197/)
Wraps bisect.bisect() in an easy to use class that supports key-functions and straight-forward search methods.
Download
Copy to clipboard1 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 99100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311 | frombisectimportbisect_left,bisect_rightclassSortedCollection(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=keykey=(lambdax:x)ifkeyisNoneelsekeydecorated=sorted((key(item),item)foriteminiterable)self._keys=[kfork,itemindecorated]self._items=[itemfork,itemindecorated]self._key=keydef_getkey(self):returnself._keydef_setkey(self,key):ifkeyisnotself._key:self.__init__(self._items,key=key)def_delkey(self):self._setkey(None)key=property(_getkey,_setkey,_delkey,'key function')defclear(self):self.__init__([],self._key)defcopy(self):returnself.__class__(self,self._key)def__len__(self):returnlen(self._items)def__getitem__(self,i):returnself._items[i]def__iter__(self):returniter(self._items)def__reversed__(self):returnreversed(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):returnself.__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)returniteminself._items[i:j]defindex(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)returnself._items[i:j].index(item)+idefcount(self,item):'Return number of occurrences of item'k=self._key(item)i=bisect_left(self._keys,k)j=bisect_right(self._keys,k)returnself._items[i:j].count(item)definsert(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)definsert_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)defremove(self,item):'Remove first occurence of item. Raise ValueError if not found'i=self.index(item)delself._keys[i]delself._items[i]deffind(self,k):'Return first item with a key == k. Raise ValueError if not found.'i=bisect_left(self._keys,k)ifi!=len(self)andself._keys[i]==k:returnself._items[i]raiseValueError('No item found with key equal to:%r'%(k,))deffind_le(self,k):'Return last item with a key <= k. Raise ValueError if not found.'i=bisect_right(self._keys,k)ifi:returnself._items[i-1]raiseValueError('No item found with key at or below:%r'%(k,))deffind_lt(self,k):'Return last item with a key < k. Raise ValueError if not found.'i=bisect_left(self._keys,k)ifi:returnself._items[i-1]raiseValueError('No item found with key below:%r'%(k,))deffind_ge(self,k):'Return first item with a key >= equal to k. Raise ValueError if not found'i=bisect_left(self._keys,k)ifi!=len(self):returnself._items[i]raiseValueError('No item found with key at or above:%r'%(k,))deffind_gt(self,k):'Return first item with a key > k. Raise ValueError if not found'i=bisect_right(self._keys,k)ifi!=len(self):returnself._items[i]raiseValueError('No item found with key above:%r'%(k,))# --------------------------- Simple demo and tests -------------------------if__name__=='__main__':defve2no(f,*args):'Convert ValueError result to -1'try:returnf(*args)exceptValueError:return-1defslow_index(seq,k):'Location of match or -1 if not found'fori,iteminenumerate(seq):ifitem==k:returnireturn-1defslow_find(seq,k):'First item with a key equal to k. -1 if not found'foriteminseq:ifitem==k:returnitemreturn-1defslow_find_le(seq,k):'Last item with a key less-than or equal to k.'foriteminreversed(seq):ifitem<=k:returnitemreturn-1defslow_find_lt(seq,k):'Last item with a key less-than k.'foriteminreversed(seq):ifitem<k:returnitemreturn-1defslow_find_ge(seq,k):'First item with a key-value greater-than or equal to k.'foriteminseq:ifitem>=k:returnitemreturn-1defslow_find_gt(seq,k):'First item with a key-value greater-than or equal to k.'foriteminseq:ifitem>k:returnitemreturn-1fromrandomimportchoicepool=[1.5,2,2.0,3,3.0,3.5,4,4.0,4.5]foriinrange(500):forninrange(6):s=[choice(pool)foriinrange(n)]sc=SortedCollection(s)s.sort()forprobeinpool:assertrepr(ve2no(sc.index,probe))==repr(slow_index(s,probe))assertrepr(ve2no(sc.find,probe))==repr(slow_find(s,probe))assertrepr(ve2no(sc.find_le,probe))==repr(slow_find_le(s,probe))assertrepr(ve2no(sc.find_lt,probe))==repr(slow_find_lt(s,probe))assertrepr(ve2no(sc.find_ge,probe))==repr(slow_find_ge(s,probe))assertrepr(ve2no(sc.find_gt,probe))==repr(slow_find_gt(s,probe))fori,iteminenumerate(s):assertrepr(item)==repr(sc[i])# test __getitem__assertiteminsc# test __contains__ and __iter__asserts.count(item)==sc.count(item)# test count()assertlen(sc)==n# test __len__assertlist(map(repr,reversed(sc)))==list(map(repr,reversed(s)))# test __reversed__assertlist(sc.copy())==list(sc)# test copy()sc.clear()# test clear()assertlen(sc)==0sd=SortedCollection('The quick Brown Fox jumped'.split(),key=str.lower)assertsd._keys==['brown','fox','jumped','quick','the']assertsd._items==['Brown','Fox','jumped','quick','The']assertsd._key==str.lowerassertrepr(sd)=="SortedCollection(['Brown', 'Fox', 'jumped', 'quick', 'The'], key=lower)"sd.key=str.upperassertsd._key==str.upperassertlen(sd)==5assertlist(reversed(sd))==['The','quick','jumped','Fox','Brown']foriteminsd:assertiteminsdfori,iteminenumerate(sd):assertitem==sd[i]sd.insert('jUmPeD')sd.insert_right('QuIcK')assertsd._keys==['BROWN','FOX','JUMPED','JUMPED','QUICK','QUICK','THE']assertsd._items==['Brown','Fox','jUmPeD','jumped','quick','QuIcK','The']assertsd.find_le('JUMPED')=='jumped',sd.find_le('JUMPED')assertsd.find_ge('JUMPED')=='jUmPeD'assertsd.find_le('GOAT')=='Fox'assertsd.find_ge('GOAT')=='jUmPeD'assertsd.find('FOX')=='Fox'assertsd[3]=='jumped'assertsd[3:5]==['jumped','quick']assertsd[-2]=='QuIcK'assertsd[-4:-2]==['jumped','quick']fori,iteminenumerate(sd):assertsd.index(item)==itry:sd.index('xyzpdq')exceptValueError:passelse:assert0,'Oops, failed to notify of missing value'sd.remove('jumped')assertlist(sd)==['Brown','Fox','jUmPeD','quick','QuIcK','The']importdoctestfromoperatorimportitemgetterprint(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
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.
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,))
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.
For a similar class with asymptotically faster inserts, see thesortedlist andsortedset classes in myblist package.
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]
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.
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).
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 thesortedcontainers available onPyPI andgithub.
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
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] | Created byRaymond HettingeronFri, 16 Apr 2010(MIT) |
| ◄ | Python recipes (4591) | ► |
| ◄ | Raymond Hettinger's recipes (97) | ► |
| ◄ | HongxuChen's Fav (39) | ► |
Required Modules
Other Information and Tasks
- Licensed under theMIT License
- Viewed 113469 times
- Revision 9(updated 13 years ago)
Accounts
Feedback & Information
ActiveState
Privacy Policy |Contact Us |Support
© 2024 ActiveState Software Inc. All rights reserved. ActiveState®, Komodo®, ActiveState Perl Dev Kit®, ActiveState Tcl Dev Kit®, ActivePerl®, ActivePython®, and ActiveTcl® are registered trademarks of ActiveState. All other marks are property of their respective owners.