intervalsets.py 2.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384
  1. # coding=utf-8
  2. #
  3. # This file is part of Hypothesis, which may be found at
  4. # https://github.com/HypothesisWorks/hypothesis-python
  5. #
  6. # Most of this work is copyright (C) 2013-2018 David R. MacIver
  7. # (david@drmaciver.com), but it contains contributions by others. See
  8. # CONTRIBUTING.rst for a full list of people who may hold copyright, and
  9. # consult the git log if you need to determine who owns an individual
  10. # contribution.
  11. #
  12. # This Source Code Form is subject to the terms of the Mozilla Public License,
  13. # v. 2.0. If a copy of the MPL was not distributed with this file, You can
  14. # obtain one at http://mozilla.org/MPL/2.0/.
  15. #
  16. # END HEADER
  17. from __future__ import division, print_function, absolute_import
  18. class IntervalSet(object):
  19. def __init__(self, intervals):
  20. self.intervals = tuple(intervals)
  21. self.offsets = [0]
  22. for u, v in self.intervals:
  23. self.offsets.append(
  24. self.offsets[-1] + v - u + 1
  25. )
  26. self.size = self.offsets.pop()
  27. def __len__(self):
  28. return self.size
  29. def __iter__(self):
  30. for u, v in self.intervals:
  31. for i in range(u, v + 1):
  32. yield i
  33. def __getitem__(self, i):
  34. if i < 0:
  35. i = self.size + i
  36. if i < 0 or i >= self.size:
  37. raise IndexError('Invalid index %d for [0, %d)' % (i, self.size))
  38. # Want j = maximal such that offsets[j] <= i
  39. j = len(self.intervals) - 1
  40. if self.offsets[j] > i:
  41. hi = j
  42. lo = 0
  43. # Invariant: offsets[lo] <= i < offsets[hi]
  44. while lo + 1 < hi:
  45. mid = (lo + hi) // 2
  46. if self.offsets[mid] <= i:
  47. lo = mid
  48. else:
  49. hi = mid
  50. j = lo
  51. t = i - self.offsets[j]
  52. u, v = self.intervals[j]
  53. r = u + t
  54. assert r <= v
  55. return r
  56. def __repr__(self):
  57. return 'IntervalSet(%r)' % (self.intervals,)
  58. def index(self, value):
  59. for offset, (u, v) in zip(self.offsets, self.intervals):
  60. if u == value:
  61. return offset
  62. elif u > value:
  63. raise ValueError('%d is not in list' % (value,))
  64. if value <= v:
  65. return offset + (value - u)
  66. raise ValueError('%d is not in list' % (value,))
  67. def index_above(self, value):
  68. for offset, (u, v) in zip(self.offsets, self.intervals):
  69. if u >= value:
  70. return offset
  71. if value <= v:
  72. return offset + (value - u)
  73. return self.size