summaryrefslogtreecommitdiff
path: root/pokemontools/interval_map.py
blob: fc53e69ca85532274fb426117ee2521e5f74878b (plain)
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
# -*- coding: utf-8 -*-

from bisect import bisect_left, bisect_right
try:
    from future_builtins import zip
except ImportError:
    pass


class IntervalMap(object):
    """
    This class maps a set of intervals to a set of values.

    >>> i = IntervalMap()
    >>> i[0:5] = "hello world"
    >>> i[6:10] = "hello cruel world"
    >>> print i[4]
    "hello world"
    """

    def __init__(self):
        """initializes an empty IntervalMap"""
        self._bounds = []
        self._items = []
        self._upperitem = None

    def __setitem__(self, _slice, _value):
        """sets an interval mapping"""
        assert isinstance(_slice, slice), 'The key must be a slice object'

        if _slice.start is None:
            start_point = -1
        else:
            start_point = bisect_left(self._bounds, _slice.start)

        if _slice.stop is None:
            end_point = -1
        else:
            end_point = bisect_left(self._bounds, _slice.stop)

        if start_point>=0:
            if start_point < len(self._bounds) and self._bounds[start_point]<_slice.start:
                start_point += 1

            if end_point>=0:
                self._bounds[start_point:end_point] = [_slice.start, _slice.stop]
                if start_point < len(self._items):
                    self._items[start_point:end_point] = [self._items[start_point], _value]
                else:
                    self._items[start_point:end_point] = [self._upperitem, _value]
            else:
                self._bounds[start_point:] = [_slice.start]
                if start_point < len(self._items):
                    self._items[start_point:] = [self._items[start_point], _value]
                else:
                    self._items[start_point:] = [self._upperitem]
                self._upperitem = _value
        else:
            if end_point>=0:
                self._bounds[:end_point] = [_slice.stop]
                self._items[:end_point] = [_value]
            else:
                self._bounds[:] = []
                self._items[:] = []
                self._upperitem = _value

    def __getitem__(self,_point):
        """gets a value from the mapping"""
        assert not isinstance(_point, slice), 'The key cannot be a slice object'

        index = bisect_right(self._bounds, _point)
        if index < len(self._bounds):
            return self._items[index]
        else:
            return self._upperitem

    def items(self):
        """returns an iterator with each item being
        ((low_bound, high_bound), value)
        these items are returned in order"""
        previous_bound = None
        for (b, v) in zip(self._bounds, self._items):
            if v is not None:
                yield (previous_bound, b), v
            previous_bound = b
        if self._upperitem is not None:
            yield (previous_bound, None), self._upperitem

    def values(self):
        """returns an iterator with each item being a stored value
        the items are returned in order"""
        for v in self._items:
            if v is not None:
                yield v
        if self._upperitem is not None:
            yield self._upperitem

    def __repr__(self):
        s = []
        for b,v in self.items():
            if v is not None:
                s.append('[%r, %r] => %r'%(
                    b[0],
                    b[1],
                    v
                ))
        return '{'+', '.join(s)+'}'