Python. Грамотно работать с любым диапазоном

Все знают, что range() в питоне используется, когда нужно что-то сделать сколько-то раз:

>>> for i in range(3, 0, -1):
...   print(i)

3
2
1

Но не все знают, что range — это коллекция (что? да!), вполне себе полноценная:

>>> seq = range(10, 100)
>>> len(seq)
90
>>> 52 in seq
True
>>> seq[10]
20

И даже так:

>>> max(seq)
99
>>> seq.index(31)
21
>>> seq.count(42)
1

И так тоже:

>>> s1 = range(0, 10, 3)
>>> s2 = range(0, 11, 3)
>>> s1 == s2
True

При этом range, в отличие от всех прочих коллекций, занимает мизерное место в памяти (48 байт), вне зависимости от того, сколько элементов в него попадают. Это потому, что хранит он только 3 атрибута: start, stop, step

>>> from pympler import asizeof
>>> seq = range(0, 100)
>>> asizeof.asizeof(seq)
48
>>> seq = range(0, 100_000)
>>> asizeof.asizeof(seq)
48
>>> seq = range(0, 100_000_000)
>>> asizeof.asizeof(seq)
48

И при этом идеальное время выполнения операций: len(), [idx], in, .index(), .count() — всё за O(1).

⌘ ⌘ ⌘

Кто-то на этом месте скажет «погодите, откуда O(1)? у списка ведь in, .index(), .count() выполняются за O(n), почему у диапазона иначе?»

Рассмотрим на примере in. Действительно: чтобы проверить, есть ли элемент в списке, придётся обходить элементы списка, пока не найдём искомый — это сложность O(n). Но в случае с диапазоном мы точно знаем первый элемент, последний элемент и шаг. Поэтому разработчики стандартной библиотеки пошли на хитрость.

Допустим, есть выражение x in range(start, stop, step). Для положительного step можно обойтись без перебора всех элементов, вот так:

def contains(range_, x):
    if x < range_.start:
        return False
    if x >= range_.stop:
        return False
    return (x - range_.start) % range_.step == 0

>>> r = range(1000, 10000, 3)
>>> contains(r, 2068)
True
>>> contains(r, 2070)
False

Проверили границы, посчитали остаток от деления, бумс, готово. Для .index() и .count() сделано аналогично, если интересно как — посмотрите исходники (осторожно, код на C):

⌘ ⌘ ⌘

Итого, получили структуру данных постоянного размера, с константным временем выполнения операций. Ну разве он не чудо, этот range?

Заметка из телеграм-канала «Oh My Py»