1424296a4e811f57ada27be68d60aeb5fcb1bb45bDoug Zongker# Copyright (C) 2014 The Android Open Source Project 2424296a4e811f57ada27be68d60aeb5fcb1bb45bDoug Zongker# 3424296a4e811f57ada27be68d60aeb5fcb1bb45bDoug Zongker# Licensed under the Apache License, Version 2.0 (the "License"); 4424296a4e811f57ada27be68d60aeb5fcb1bb45bDoug Zongker# you may not use this file except in compliance with the License. 5424296a4e811f57ada27be68d60aeb5fcb1bb45bDoug Zongker# You may obtain a copy of the License at 6424296a4e811f57ada27be68d60aeb5fcb1bb45bDoug Zongker# 7424296a4e811f57ada27be68d60aeb5fcb1bb45bDoug Zongker# http://www.apache.org/licenses/LICENSE-2.0 8424296a4e811f57ada27be68d60aeb5fcb1bb45bDoug Zongker# 9424296a4e811f57ada27be68d60aeb5fcb1bb45bDoug Zongker# Unless required by applicable law or agreed to in writing, software 10424296a4e811f57ada27be68d60aeb5fcb1bb45bDoug Zongker# distributed under the License is distributed on an "AS IS" BASIS, 11424296a4e811f57ada27be68d60aeb5fcb1bb45bDoug Zongker# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12424296a4e811f57ada27be68d60aeb5fcb1bb45bDoug Zongker# See the License for the specific language governing permissions and 13424296a4e811f57ada27be68d60aeb5fcb1bb45bDoug Zongker# limitations under the License. 14424296a4e811f57ada27be68d60aeb5fcb1bb45bDoug Zongker 15fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongkerfrom __future__ import print_function 16fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongkerimport heapq 17fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongkerimport itertools 18fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker 19fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker__all__ = ["RangeSet"] 20fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker 21fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongkerclass RangeSet(object): 22fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker """A RangeSet represents a set of nonoverlapping ranges on the 23fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker integers (ie, a set of integers, but efficient when the set contains 24fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker lots of runs.""" 25fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker 26fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker def __init__(self, data=None): 27fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker if data: 28fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker self.data = tuple(self._remove_pairs(data)) 29fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker else: 30fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker self.data = () 31fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker 32fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker def __iter__(self): 33fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker for i in range(0, len(self.data), 2): 34fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker yield self.data[i:i+2] 35fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker 36fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker def __eq__(self, other): 37fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker return self.data == other.data 38fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker def __ne__(self, other): 39fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker return self.data != other.data 40fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker def __nonzero__(self): 41fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker return bool(self.data) 42fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker 43fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker def __str__(self): 44fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker if not self.data: 45fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker return "empty" 46fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker else: 47fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker return self.to_string() 48fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker 49fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker @classmethod 50fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker def parse(cls, text): 51fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker """Parse a text string consisting of a space-separated list of 52fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker blocks and ranges, eg "10-20 30 35-40". Ranges are interpreted to 53fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker include both their ends (so the above example represents 18 54fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker individual blocks. Returns a RangeSet object. 55fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker 56fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker If the input has all its blocks in increasing order, then returned 57fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker RangeSet will have an extra attribute 'monotonic' that is set to 58fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker True. For example the input "10-20 30" is monotonic, but the input 59fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker "15-20 30 10-14" is not, even though they represent the same set 60fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker of blocks (and the two RangeSets will compare equal with ==). 61fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker """ 62fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker 63fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker data = [] 64fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker last = -1 65fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker monotonic = True 66fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker for p in text.split(): 67fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker if "-" in p: 68fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker s, e = p.split("-") 69fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker data.append(int(s)) 70fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker data.append(int(e)+1) 71fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker if last <= s <= e: 72fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker last = e 73fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker else: 74fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker monotonic = False 75fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker else: 76fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker s = int(p) 77fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker data.append(s) 78fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker data.append(s+1) 79fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker if last <= s: 80fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker last = s+1 81fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker else: 82fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker monotonic = True 83fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker data.sort() 84fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker r = RangeSet(cls._remove_pairs(data)) 85fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker r.monotonic = monotonic 86fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker return r 87fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker 88fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker @staticmethod 89fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker def _remove_pairs(source): 90fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker last = None 91fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker for i in source: 92fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker if i == last: 93fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker last = None 94fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker else: 95fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker if last is not None: 96fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker yield last 97fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker last = i 98fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker if last is not None: 99fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker yield last 100fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker 101fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker def to_string(self): 102fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker out = [] 103fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker for i in range(0, len(self.data), 2): 104fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker s, e = self.data[i:i+2] 105fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker if e == s+1: 106fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker out.append(str(s)) 107fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker else: 108fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker out.append(str(s) + "-" + str(e-1)) 109fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker return " ".join(out) 110fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker 111fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker def to_string_raw(self): 112fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker return str(len(self.data)) + "," + ",".join(str(i) for i in self.data) 113fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker 114fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker def union(self, other): 115fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker """Return a new RangeSet representing the union of this RangeSet 116fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker with the argument.""" 117fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker out = [] 118fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker z = 0 119fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker for p, d in heapq.merge(zip(self.data, itertools.cycle((+1, -1))), 120fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker zip(other.data, itertools.cycle((+1, -1)))): 121fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker if (z == 0 and d == 1) or (z == 1 and d == -1): 122fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker out.append(p) 123fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker z += d 124fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker return RangeSet(data=out) 125fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker 126fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker def intersect(self, other): 127fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker """Return a new RangeSet representing the intersection of this 128fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker RangeSet with the argument.""" 129fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker out = [] 130fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker z = 0 131fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker for p, d in heapq.merge(zip(self.data, itertools.cycle((+1, -1))), 132fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker zip(other.data, itertools.cycle((+1, -1)))): 133fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker if (z == 1 and d == 1) or (z == 2 and d == -1): 134fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker out.append(p) 135fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker z += d 136fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker return RangeSet(data=out) 137fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker 138fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker def subtract(self, other): 139fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker """Return a new RangeSet representing subtracting the argument 140fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker from this RangeSet.""" 141fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker 142fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker out = [] 143fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker z = 0 144fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker for p, d in heapq.merge(zip(self.data, itertools.cycle((+1, -1))), 145fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker zip(other.data, itertools.cycle((-1, +1)))): 146fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker if (z == 0 and d == 1) or (z == 1 and d == -1): 147fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker out.append(p) 148fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker z += d 149fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker return RangeSet(data=out) 150fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker 151fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker def overlaps(self, other): 152fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker """Returns true if the argument has a nonempty overlap with this 153fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker RangeSet.""" 154fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker 155fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker # This is like intersect, but we can stop as soon as we discover the 156fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker # output is going to be nonempty. 157fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker z = 0 158fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker for p, d in heapq.merge(zip(self.data, itertools.cycle((+1, -1))), 159fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker zip(other.data, itertools.cycle((+1, -1)))): 160fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker if (z == 1 and d == 1) or (z == 2 and d == -1): 161fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker return True 162fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker z += d 163fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker return False 164fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker 165fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker def size(self): 166fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker """Returns the total size of the RangeSet (ie, how many integers 167fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker are in the set).""" 168fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker 169fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker total = 0 170fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker for i, p in enumerate(self.data): 171fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker if i % 2: 172fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker total += p 173fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker else: 174fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker total -= p 175fc44a515d46e6f4d5eaa0d32659b1cf3b9492305Doug Zongker return total 176