intranges.py 1.9 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758
  1. """
  2. Given a list of integers, made up of (hopefully) a small number of long runs
  3. of consecutive integers, compute a representation of the form
  4. ((start1, end1), (start2, end2) ...). Then answer the question "was x present
  5. in the original list?" in time O(log(# runs)).
  6. """
  7. import bisect
  8. from typing import List, Tuple
  9. def intranges_from_list(list_):
  10. # type: (List[int]) -> Tuple[int, ...]
  11. """Represent a list of integers as a sequence of ranges:
  12. ((start_0, end_0), (start_1, end_1), ...), such that the original
  13. integers are exactly those x such that start_i <= x < end_i for some i.
  14. Ranges are encoded as single integers (start << 32 | end), not as tuples.
  15. """
  16. sorted_list = sorted(list_)
  17. ranges = []
  18. last_write = -1
  19. for i in range(len(sorted_list)):
  20. if i+1 < len(sorted_list):
  21. if sorted_list[i] == sorted_list[i+1]-1:
  22. continue
  23. current_range = sorted_list[last_write+1:i+1]
  24. ranges.append(_encode_range(current_range[0], current_range[-1] + 1))
  25. last_write = i
  26. return tuple(ranges)
  27. def _encode_range(start, end):
  28. # type: (int, int) -> int
  29. return (start << 32) | end
  30. def _decode_range(r):
  31. # type: (int) -> Tuple[int, int]
  32. return (r >> 32), (r & ((1 << 32) - 1))
  33. def intranges_contain(int_, ranges):
  34. # type: (int, Tuple[int, ...]) -> bool
  35. """Determine if `int_` falls into one of the ranges in `ranges`."""
  36. tuple_ = _encode_range(int_, 0)
  37. pos = bisect.bisect_left(ranges, tuple_)
  38. # we could be immediately ahead of a tuple (start, end)
  39. # with start < int_ <= end
  40. if pos > 0:
  41. left, right = _decode_range(ranges[pos-1])
  42. if left <= int_ < right:
  43. return True
  44. # or we could be immediately behind a tuple (int_, end)
  45. if pos < len(ranges):
  46. left, _ = _decode_range(ranges[pos])
  47. if left == int_:
  48. return True
  49. return False