• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1# Note: Taken from https://code.activestate.com/recipes/577197-sortedcollection/.
2#
3# Copyright 2010 Raymond Hettinger
4#
5# Permission is hereby granted, free of charge, to any person obtaining a copy of
6# this software and associated documentation files (the "Software"), to deal in the
7# Software without restriction, including without limitation the rights to use, copy,
8# modify, merge, publish, distribute, sublicense, and/or sell copies of the Software,
9# and to permit persons to whom the Software is furnished to do so, subject to the
10# following conditions:
11#
12# The above copyright notice and this permission notice shall be included in all copies
13# or substantial portions of the Software.
14#
15# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED,
16# INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR
17# PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE
18# FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR
19# OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
20# DEALINGS IN THE SOFTWARE.
21
22from bisect import bisect_left, bisect_right
23
24class SortedCollection(object):
25    def __init__(self, iterable=(), key=None):
26        self._given_key = key
27        key = (lambda x: x) if key is None else key
28        decorated = sorted((key(item), item) for item in iterable)
29        self._keys = [k for k, item in decorated]
30        self._items = [item for k, item in decorated]
31        self._key = key
32
33    def _getkey(self):
34        return self._key
35
36    def _setkey(self, key):
37        if key is not self._key:
38            self.__init__(self._items, key=key)
39
40    def _delkey(self):
41        self._setkey(None)
42
43    key = property(_getkey, _setkey, _delkey, 'key function')
44
45    def clear(self):
46        self.__init__([], self._key)
47
48    def copy(self):
49        return self.__class__(self, self._key)
50
51    def __len__(self):
52        return len(self._items)
53
54    def __getitem__(self, i):
55        return self._items[i]
56
57    def __iter__(self):
58        return iter(self._items)
59
60    def __reversed__(self):
61        return reversed(self._items)
62
63    def __repr__(self):
64        return '%s(%r, key=%s)' % (
65            self.__class__.__name__,
66            self._items,
67            getattr(self._given_key, '__name__', repr(self._given_key))
68        )
69
70    def __reduce__(self):
71        return self.__class__, (self._items, self._given_key)
72
73    def __contains__(self, item):
74        k = self._key(item)
75        i = bisect_left(self._keys, k)
76        j = bisect_right(self._keys, k)
77        return item in self._items[i:j]
78
79    def index(self, item):
80        'Find the position of an item.  Raise ValueError if not found.'
81        k = self._key(item)
82        i = bisect_left(self._keys, k)
83        j = bisect_right(self._keys, k)
84        return self._items[i:j].index(item) + i
85
86    def count(self, item):
87        'Return number of occurrences of item'
88        k = self._key(item)
89        i = bisect_left(self._keys, k)
90        j = bisect_right(self._keys, k)
91        return self._items[i:j].count(item)
92
93    def insert(self, item):
94        'Insert a new item.  If equal keys are found, add to the left'
95        k = self._key(item)
96        i = bisect_left(self._keys, k)
97        self._keys.insert(i, k)
98        self._items.insert(i, item)
99
100    def insert_right(self, item):
101        'Insert a new item.  If equal keys are found, add to the right'
102        k = self._key(item)
103        i = bisect_right(self._keys, k)
104        self._keys.insert(i, k)
105        self._items.insert(i, item)
106
107    def remove(self, item):
108        'Remove first occurence of item.  Raise ValueError if not found'
109        i = self.index(item)
110        del self._keys[i]
111        del self._items[i]
112
113    def find(self, k):
114        'Return first item with a key == k.  Raise ValueError if not found.'
115        i = bisect_left(self._keys, k)
116        if i != len(self) and self._keys[i] == k:
117            return self._items[i]
118        raise ValueError('No item found with key equal to: %r' % (k,))
119
120    def find_le(self, k):
121        'Return last item with a key <= k.  Raise ValueError if not found.'
122        i = bisect_right(self._keys, k)
123        if i:
124            return self._items[i-1]
125        raise ValueError('No item found with key at or below: %r' % (k,))
126
127    def find_lt(self, k):
128        'Return last item with a key < k.  Raise ValueError if not found.'
129        i = bisect_left(self._keys, k)
130        if i:
131            return self._items[i-1]
132        raise ValueError('No item found with key below: %r' % (k,))
133
134    def find_ge(self, k):
135        'Return first item with a key >= equal to k.  Raise ValueError if not found'
136        i = bisect_left(self._keys, k)
137        if i != len(self):
138            return self._items[i]
139        raise ValueError('No item found with key at or above: %r' % (k,))
140
141    def find_gt(self, k):
142        'Return first item with a key > k.  Raise ValueError if not found'
143        i = bisect_right(self._keys, k)
144        if i != len(self):
145            return self._items[i]
146        raise ValueError('No item found with key above: %r' % (k,))
147