bisect.py 3.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118
  1. """Bisection algorithms."""
  2. def insort_right(a, x, lo=0, hi=None, *, key=None):
  3. """Insert item x in list a, and keep it sorted assuming a is sorted.
  4. If x is already in a, insert it to the right of the rightmost x.
  5. Optional args lo (default 0) and hi (default len(a)) bound the
  6. slice of a to be searched.
  7. A custom key function can be supplied to customize the sort order.
  8. """
  9. if key is None:
  10. lo = bisect_right(a, x, lo, hi)
  11. else:
  12. lo = bisect_right(a, key(x), lo, hi, key=key)
  13. a.insert(lo, x)
  14. def bisect_right(a, x, lo=0, hi=None, *, key=None):
  15. """Return the index where to insert item x in list a, assuming a is sorted.
  16. The return value i is such that all e in a[:i] have e <= x, and all e in
  17. a[i:] have e > x. So if x already appears in the list, a.insert(i, x) will
  18. insert just after the rightmost x already there.
  19. Optional args lo (default 0) and hi (default len(a)) bound the
  20. slice of a to be searched.
  21. A custom key function can be supplied to customize the sort order.
  22. """
  23. if lo < 0:
  24. raise ValueError('lo must be non-negative')
  25. if hi is None:
  26. hi = len(a)
  27. # Note, the comparison uses "<" to match the
  28. # __lt__() logic in list.sort() and in heapq.
  29. if key is None:
  30. while lo < hi:
  31. mid = (lo + hi) // 2
  32. if x < a[mid]:
  33. hi = mid
  34. else:
  35. lo = mid + 1
  36. else:
  37. while lo < hi:
  38. mid = (lo + hi) // 2
  39. if x < key(a[mid]):
  40. hi = mid
  41. else:
  42. lo = mid + 1
  43. return lo
  44. def insort_left(a, x, lo=0, hi=None, *, key=None):
  45. """Insert item x in list a, and keep it sorted assuming a is sorted.
  46. If x is already in a, insert it to the left of the leftmost x.
  47. Optional args lo (default 0) and hi (default len(a)) bound the
  48. slice of a to be searched.
  49. A custom key function can be supplied to customize the sort order.
  50. """
  51. if key is None:
  52. lo = bisect_left(a, x, lo, hi)
  53. else:
  54. lo = bisect_left(a, key(x), lo, hi, key=key)
  55. a.insert(lo, x)
  56. def bisect_left(a, x, lo=0, hi=None, *, key=None):
  57. """Return the index where to insert item x in list a, assuming a is sorted.
  58. The return value i is such that all e in a[:i] have e < x, and all e in
  59. a[i:] have e >= x. So if x already appears in the list, a.insert(i, x) will
  60. insert just before the leftmost x already there.
  61. Optional args lo (default 0) and hi (default len(a)) bound the
  62. slice of a to be searched.
  63. A custom key function can be supplied to customize the sort order.
  64. """
  65. if lo < 0:
  66. raise ValueError('lo must be non-negative')
  67. if hi is None:
  68. hi = len(a)
  69. # Note, the comparison uses "<" to match the
  70. # __lt__() logic in list.sort() and in heapq.
  71. if key is None:
  72. while lo < hi:
  73. mid = (lo + hi) // 2
  74. if a[mid] < x:
  75. lo = mid + 1
  76. else:
  77. hi = mid
  78. else:
  79. while lo < hi:
  80. mid = (lo + hi) // 2
  81. if key(a[mid]) < x:
  82. lo = mid + 1
  83. else:
  84. hi = mid
  85. return lo
  86. # Overwrite above definitions with a fast C implementation
  87. try:
  88. from _bisect import *
  89. except ImportError:
  90. pass
  91. # Create aliases
  92. bisect = bisect_right
  93. insort = insort_right