Итераторы и генераторы

В чем разница между итератором и генератором? Этот вопрос можно часто услышать на собеседованиях.

Итератор – более общая концепция, чем генератор.

Итератор – это интерфейс доступа к элементам коллекций и потоков данных. Он требует реализации единственного метода – «дай мне следующий элемент». Если вы пишите свой итератор на Python 3 вам нужно реализовать в классе метод __next__. Если элементы исчерпаны итератор возбудит исключение StopIteration.

📎 Пример. Итератор счетчик – выдает числа от low до high:

class Counter:
    def __init__(self, low, high):
        self.current = low
        self.high = high
    def __iter__(self):
        return self
    def __next__(self): 
        if self.current > self.high:
            raise StopIteration
        else:
            self.current += 1
            return self.current - 1

Генератор – это итератор

Генератор – это итератор, но не наоборот. Не любой итератор является генератором.

Есть два способа получить генератор:

📎 1. Генераторное выражение (что-то типа list comprehension, но возвращает генератор, а не список). Используются круглые скобки:

>>> g = (2 * i for i in range(5))
>>> type(g)
<class 'generator'>
>>> next(g)
0
>>> next(g)
2

📎 2. Генераторные функции. Это функции, где есть хотя бы одно выражение yield. Когда мы запускаем генератор, функция выполняет до первого выражения yield. То, что мы передали в yield будет возвращено наружу. Генератор при этом встанет «на паузу» до следующей итерации. При следующей итерации выполнение генератора продолжится до очередного yield.

Генераторы можно прочитать только 1 раз, потому что обычно генераторы не хранят значения в памяти, а генерируют их налету (отсюда и название).

Пример. Генератор чисел Фибоначчи (бесконечный):

def fib():
    a, b = 0, 1
    while 1:
        yield a
        a, b = b, a + b

>>> fib_g = fib()
>>> next(fib_g)
0
>>> next(fib_g)
1
>>> next(fib_g)
1
>>> next(fib_g)
2
>>> next(fib_g)
3
>>> next(fib_g)
5

Вызвав генераторную функцию fib() мы получили генератор. Затем мы итерируем этот генератор функцией next().

Передача данных в генератор

У генераторов есть дополнительные методы, которые позволяют передавать внутрь генератора данные или возбуждать внутри него исключения. Это еще одно отличие от простых итераторов.

send() – отправить данные в генератор. Переданное значение вернется из той конструкции yield, на которой возникла последняя пауза генератора. При этом генератор будет прокручен на один шаг, как если бы мы вызвали next:

val = yield i  # генератор вернет i, но внутри получит val из аргумента метода send

Пример. Этот генератор просто выдает числа от 0 и далее, при этом печатает в поток вывода все, что мы ему отправляем.

def my_gen():
    i = 0
    while True:
        val = yield i
        print('Got inside generator:', val)
        i += 1

>>> g = my_gen()
>>> next(g)
0
>>> g.send("hello")
Got inside generator: hello
1
>>> g.send("world")
Got inside generator: world
2

Обратите внимание, что первый раз нельзя посылать в генератор данные, пока мы не прокрутили его до первого yield. Нужно либо взывать next(g) или g.send(None) – это одно и тоже.

Не будет ошибкой отправлять данные генератору, который не получает их (нет использования значения конструкции yield). Например, нашему генератору fib() можно отравить все, что угодно, он просто проигнорирует.

throw() – бросить исключение внутри генератора. Исключение будет возбуждено из того выражение yield, где генератор последний раз остановился.

>>> g = my_gen()   # my_gen из прошлого примера

>>> g.throw(TypeError, 'my error')
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "<stdin>", line 1, in my_gen
TypeError: my error

close() – закрыть генератор. Бросает внутри генератора особое исключение GeneratorExit. Это исключение, даже если оно не обработано, не распространится в код, вызвавший close(). Но, если мы поймали это исключение внутри генератора, то после закрытия генератора нельзя уже делать yield, рискуя получить RuntimeError. Остальные виды исключений будут распространяться из генератора в код, его вызывающий. Попытка итерировать закрытый итератор приведет к исключению StopIteration (закрытый генератор – пустой итератор).

>>> g = my_gen()
>>> next(g)
0
>>> next(g)
Got inside generator: None
1
>>> g.close()
>>> next(g)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
StopIteration

Бонус

Как взять из итератора (в том числе из генератора) N первых значений?

Можно, конечно, написать свою функцию. Но зачем, если она уже есть в стандартном модуле itertools. Этот модуль содержит множество вспомогательных функций для работы с итераторами. Нам понадобится itertools.islice. Первый аргумент – итератор (ну или генератор), остальные три – как в range.

>>> list(itertools.islice(fib(), 10))
[0, 1, 1, 2, 3, 5, 8, 13, 21, 34]

>>> list(itertools.islice(fib(), 10, 20, 2))
[55, 144, 377, 987, 2584]

В первом примере мы передаем в функцию itertools.islice наш генератор чисел Фибоначчи и число чисел, которые надо вычислить (в нашем случае – 10).

Мы также применяем функцию list, чтобы посмотреть список значений, потому что itertools.islice возвращает не спикок, а именно новый итератор, в котором будут только интересные нам значений из исходного итератора.

Во втором примеры аргументов 4 штуки. В этом случае второй аргумент – начальный номер = 10, третий – конечный номер = 20 – (не включительно), и четвертый – шаг = 2. (Очень похоже на range, не так ли?)

Специально для канала @pyway.

Python: is

Новички часто путаются в конструкциях is и ==. Давайте разберемся, что к чему.

Сразу к сути: == (и его антагонист !=) применяются для проверки равенства (неравенства) значения двух объектов. Значение, это непосредственно то, что лежит в переменной. Значение числа 323235 – собственно число 323235. Тавтология. Но на примерах станет яснее.

Оператор is (и его антагонист is not) применяются проверки равенства (неравенства) ссылок на объект. Сразу отметим то, что на значение (допустим 323235) может быть копировано и храниться в разных местах (в разных объектах в памяти).

>>> x = 323235
>>> y = 323235
>>> x == y
True
>>> x is y
False

Видите, значение переменных равны по значению, но они ссылаются на разные объекты. Я не случайно взял большое число 323235. Дело в том, что в целях оптимизации интерпретатор Python при старте создает некоторые количество часто-используемых констант (от -5 до 256 включительно).

Следите внимательно за ловкостью рук:

>>> x = 256
>>> y = 256
>>> x is y
True
>>> x = 257
>>> y = 257
>>> x is y
False
>>> x = -5
>>> y = -5
>>> x is y
True
>>> x = -6
>>> y = -6
>>> x is y
False 

Поэтому новички часто совершают ошибку, считая, что писать == – это как-то не Python-way, а is – Python-way. Это ошибочное предположение может быть раскрыто не сразу.

Python старается кэшировать и переиспользовать строковые значения. Поэтому весьма вероятно, что переменные, содержащие одинаковые строки, будут содержать ссылки на одинаковые объекты. Но это не факт! Смотрите последний пример:

>>> x = "hello"
>>> y = "hello"
>>> x is y
True
>>> x = "hel" + "lo"
>>> y = "hello"
>>> x is y
True
>>> a = "hel"
>>> b = "lo"
>>> x = a + b
>>> y = "hello"
>>> x == y
True
>>> x is y
False

Мы составили строку из двух частей и она попала в другой объект. Python не догадался (и правильно) поискать ее в существующих строках.

Суть is (id)

В Python есть встроенная функция id. Она возвращает идентификатор объекта – некоторое число. Гарантируется, что оно будет различно для различных объектах в пределах одного интерпретатора. В реализации CPython – это просто адрес объекта в памяти интерпретатора.

Так вот:

a is b

Это тоже самое, что:

id(a) == id(b)

И все! Пример для проверки:

>>> x = 10.40
>>> y = 10.40
>>> x is y
False
>>> x == y
True

>>> id(x)
4453475504
>>> id(y)
4453475600
>>> id(x) == id(y)
False

>>> x = y
>>> x is y
True
>>> id(x)
4453475600
>>> id(y)
4453475600

Значения переменных равны, но их id – разные, и is выдает False. Как только мы к x привязали y, то ссылки стали совпадать.

Для чего можно применять is?

Если мы точно знаем уверены, что хотим проверять именно равенство ссылок на объекты (один ли это объект в памяти или разные).

Еще можно применять is для сравнения с None. None – это встроенная константа и двух None быть не может.

>>> x is None
False
>>> x = None
>>> x is None
True

Также для Ellipsis:

>>> ... is Ellipsis
True
>>> x = ...
>>> y = ...
>>> x is y
True

Я не рекомендую применять is для True и False.

Потому что короче писать if x:, чем if x is True:.

Можно применять is для сравнения типов с осторожностью (без учета наследования, т. е. проверка на точное совпадение типов):

>>> x = 10.5
>>> type(x) is float
True

С наследованием может быть конфуз:

>>> class Foo: ...
...
>>> class Bar(Foo): ...
...
>>> f = Foo()
>>> b = Bar()
>>> type(f) is Foo
True
>>> type(b) is Bar
True
>>> type(b) is Foo
False
>>> isinstance(b, Foo)
True

Не смотря на то, что Bar – наследник Foo, типы переменных foo и bar не совпадают. Если нам важно учесть наcледование, то пишите isinstance.

Нюанс: is not против is (not)

Важно знать, что is not – это один целый оператор, аналогичный id(x) != id(y). А в конструкции x is (not y) – у нас сначала будет логическое отрицание y, а потом просто оператор is.

Пример уловки:

>>> x = 10
>>> x is not None
True
>>> x is (not None)
False

Сравнение пользовательских классов

Далее речь пойдет об обычных == и !=. Можно определить магический метод __eq__, который обеспечит поведение при сравнении классов. Если он не реализован, то объекты будет сравниваться по ссылкам (как при is).

>>> class Baz: ...
...
>>> x = Baz()
>>> y = Baz()
>>> x == y
False
>>> x = y
>>> x == y
True

Если он реализован, то будет вызван метод __eq__ для левого операнда.

class Foo:
 def __init__(self, x):
  self.x = x
 def __eq__(self, other):
  print('Foo __eq__ {} and {}'.format(self, other))
  return self.x == other.x

>>> x = Foo(5)
>>> y = Foo(5)
>>> x == y
Foo __eq__ <__main__.Foo object at 0x109e9c048> and <__main__.Foo object at 0x109e8a5c0>
True

Метод __ne__ отвечает за реализацию !=. По умолчанию он вызывает not x.__eq__(y). Но рекомендуется реализовывать их оба вручную, чтобы поведение сравнения было согласовано и явно.

Вопрос к размышлению: что будет если мы сравним объекты разных классов, причем оба класса реализуют __eq__?

Что будет, если мы реализуем __ne__, но не реализуем __eq__?

А еще есть метод __cmp__. Это уже выходит за рамки статьи про is. Почитайте самостоятельно…

Специально для канала @pyway.

🖨 Все о print

print – одна из первый функций, с которой знакомятся новички, изучающие Python.

>>> print("Hello world!")
Hello world!

Однако, print умеет больше, чем просто печатать через пробел. Взглянем на сигнатуру:

print(*objects, sep=' ', end='\n', file=sys.stdout, flush=False)

Именованные параметры управляют поведением функции.

Во-первых, обратим внимание на параметр sep – это строка-разделитель, она будет вставлена между каждым из неименованных параметров переданным в print (по умолчанию – одиночный пробел).

>>> print(1, 2, 3, 4, sep=" and ")
1 and 2 and 3 and 4

Если передать sep=’\n’, то каждый параметр будет выведен с новой строки.

Во-вторых, параметр end – будет выведен после печати последнего аргумента. По умолчанию, перевод строки. Например, иногда требуется печатать в одну строку: тогда передим end=»:

>>> print("Same", end=''); print("Line")
SameLine

Параметр file указывает файл или поток, в который будет выводится информация, по умолчанию, это стандартный вывод sys.stdout. Можно печатать в файл:

>>> f = open('1.txt', 'w')
>>> print("hello", file=f)
>>> f.close()

Параметр flush=True заставляет систему немедленно сбросить содержимое буфера в поток вывода. Изредка данные застревают в буфере, и этим мы их проталкиваем. Например, при end=», «1» появится только через 5 секунд сразу с «2».

>>> print('1', end=''); time.sleep(5); print('2')

Исправим это:

>>> print('1', end='', flush=True); time.sleep(5); print('2')

Наконец, не забудем, что бывает удобно печатать коллекции, передав их в print со звездочкой.

>>> arr = ["Apple", "pear", "orange"]
>>> print(*arr, sep=", ")
Apple, pear, orange

Индексирование в Python

Положительные и отрицательные индексы

Допустим у нас есть список или кортеж.

x = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
t = (0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10)

Без потери общности будем работать только со списком х (с кортежем t – тоже самое).

Легко получить i-тый элемент этого списка по индексу.

Внимание! Индексы в Python считаются с нуля (0), как в С++ и Java.

>>> x[0]
0
>>> x[7]
7
>>> x[11]
Traceback (most recent call last):
File "<stdin>", line 1, in <module>
IndexError: list index out of range

В последней строке мы вылезли за пределы (у нас в списке последний индекс – 10) и получили исключение IndexError.

Но что будет, если мы обратимся к элементу с отрицательным индексом? В С++ такой операцией вы бы прострелили себе ногу. А в Python? IndexError? Нет!

>>> x[-1]
10
>>> x[-2]
9
>>> x[-10]
1
>>> x[-11]
0

Это совершенно легально. Мы просто получаем элементы не с начала списка, а с конца (-i-тый элемент).
x[-1] – последний элемент.
x[-2] – предпоследний элемент.

Это аналогично конструкции x[len(x)-i]:

>>> x[len(x)-1]
10

Обратите внимание, что начальный (слева) элемент в отрицательной нотации имеет индекс -11.

Срезы

Срезы, они же slices, позволяют вам получить какую-то часть списка или кортежа.

Форма x[start:end] даст элементы от индекса start (включительно) до end (не включая end). Если не указать start – мы начнем с 0-го элемента, если не указать end – то закончим последним элементом (включительно). Соотвественно, x[:] это тоже самое, что и просто x.

>>> x = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
>>> x[2:8]
[2, 3, 4, 5, 6, 7]
>>> x[:8]
[0, 1, 2, 3, 4, 5, 6, 7]
>>> x[2:]
[2, 3, 4, 5, 6, 7, 8, 9, 10]
>>> x[:]
[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]

Если end <= start, получим пустой список.

>>> x[5:3]
[]

Аналогично мы можем получать срезы с отчетом от конца списка с помощью отрицательных индексов.

>>> a = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
>>> a[-4:-2]
[7, 8]

В этом случае также start < end, иначе будет пустой список.

Форма x[start:end:step] даст элементы от индекса start (включительно) до end (не включая end), в шагом step. Если step равен 1, то эта форма аналогична предыдущей рассмотренной x[start:end].

>>> x[::2]
[0, 2, 4, 6, 8, 10]
>>> x[::3]
[0, 3, 6, 9]
>>> x[2:8:2]
[2, 4, 6]

x[::2] – каждый второй элемент, а x[::3] – каждый третий. 

Отрицательный шаг вернет нам элементы в обратном порядке:

>>> x[::-1]
[10, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0]

# как если бы:
>>> list(reversed(x))
[10, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0]

# в обратном порядке с шагом 2
>>> x[::-2]
[10, 8, 6, 4, 2, 0]

Запись в список по срезу

Можно присвоить части списка, отобранной срезом, некоторый другой список, причем размер среза не обязан равняться размеру присваемого списка.

Если размеры равны (в примере два элемента в срезе и два элемента во втором списке) – происходит замена элементов.

>>> a = [1,2,3,4,5]
>>> a[1:3] = [22, 33]
>>> a
[1, 22, 33, 4, 5]

Если они не равны по размеру, то в результате список расширяется или сжимается.

>>> a = [1, 2, 3, 4, 5]
# размер среза = 1 элемент, а вставляем два (массив расширился)
>>> a[2:3] = [0, 0]
>>> a
[1, 2, 0, 0, 4, 5]

# тут вообще пустой размер среза = вставка подсписка по индексу 1
>>> a[1:1] = [8, 9]
>>> a
[1, 8, 9, 2, 0, 0, 4, 5]

# начиная с элемента 1 и кончая предпоследним элементом мы уберем (присвоив пустой список)
>>> a[1:-1] = []
>>> a
[1, 5]

Именованные срезы

Можно заранее создавать срезы с какими-то параметрами без привязки к списку или кортежу встроенной функцией slice. А потом применить этот срез к какому-то списку.

>>> a = [0, 1, 2, 3, 4, 5]
>>> LASTTHREE = slice(-3, None)
>>> LASTTHREE
slice(-3, None, None)
>>> a[LASTTHREE]
[3, 4, 5]

Вместо пустых мест для start, end или step здесь мы пишем None.

В заключение к этому разделу хочу сказать, что срезы списков возвращают списки, срезы кортежей – кортежи.

Индексирование своих объектов

В конце концов, мы можете определить самостоятельно поведение оператор индексации [], определив для своего класса магические методы __getitem__, __setitem__ и __delitem__. Первый вызывается при получении значения по индекса (или индексам), второй – если мы попытаемся нашему объекту что-то присвоить по индексу. А третий – если мы будет пытаться делать del по индексу. Необязательно реализовывать их все. Можно только один, например:

# при чтении по индексу из этого класса, мы получим удвоенных индекс
class MyClass:
    def __getitem__(self, key):
       return key * 2

myobj = MyClass()
myobj[3]  # вернет 6
myobj["privet!"] # приколись, будет: 'privet!privet!'

В качестве ключей можно использовать не только целые числа, но и строки или любые другие значения, в том числе slice и Ellipsis. Как вы будете обрабатывать их – решать вам. Естественно, логика, описанная в предыдущих разделах, здесь будет только в том случае, если вы ее сами так реализуете.

Пример. Экземпляр этого класса возвращаем нам список из целых чисел по индексу в виде срезу. Этакий бесконечный массив целых чисел, который почти не занимает памяти.

class IntegerNumbers:
  def __getitem__(self, key):
    if isinstance(key, int):
      return key
    elif isinstance(key, slice):
      return list(range(key.start, key.stop, key.step))
    else:
      raise ValueError

ints = IntegerNumbers() 
print(ints[10])  # 10
print(ints[1:10:2]) # [1, 3, 5, 7, 9]
print(ints["wwdwd"]) # так нельзя

Можно иметь несколько индексов. Ниже мы суммируем все значения индексов.

class MultiIndex:
  def __getitem__(self, keys): 
    # все индексы (если их 2 и больше попадут) в keys с типом tuple
    return sum(keys)  # просуммируем их

prod = MultiIndex()
print(prod[10, 20])  # напечает 30

Удачи в программировании и жизни!

Специально для канала PyWay.

Троеточие

Троеточие. Оно же многоточие, если по правилам, как утверждала наша учительница русского языка в школе и ставила нам двойки. Что тут сказать? Для некоторых все, что больше двух – уже много.

В питоне есть такая вещь:

...

Да это три простые точки подряд без пробелов.

Зачем она нужна и что это такое? Это Ellipsis, по-русски – оно самое «…точие». По-научному – это литерал встроенной константы (Ellipsis). Типа есть True, False, None, … и Ellipsis.

>>> ...
Ellipsis
>>> type(...)
<class 'ellipsis'>
>>> bool(...)
True

Т. е. троеточие это типа такое значение, которое вычисляется в константу Ellipsis, которая сама по себе равна себе же, имеет класс ellipsis и может быть приведена к bool=True, если надо.

Зачем она нужна? Да черт ее знает… Не припомню мест, где она используется в стандартной библиотеке.

Зато можно креативно использовать ее там, где у нас пустое тело класса или оператора (вместо pass).

class Test:
    pass

# красивее будет
# типа пока здесь пусто, но продолжение следует. загадочный и многозначительный код

class Test:
    ...

# или

if x == 5:
   ...  # я сюда еще чего-нить допишу потом
else: 
   print('not 5')

Можно троеточие передавать как значение аргумента ф-ции. Пример:

def foo(x):
    if x is Ellipsis:
        print('mda.....')
    else:
        print('x =', x)

foo(6)    # напечатает: x = 6
foo(...)  # напечатает: mda.....

Реальные применения троеточия встречаются в библиотеке для математических вычислений numpy.

Там он позволяет выбрать из многомерного массива подмассив, используя полную выборку по тем измерениям, которые не указаны конкретно. Сам не понял, что сказал, тут надо быть математиком!

В примере ниже у нас 4-х мерный массив. Мы фиксируем первый и последний индексы, а 2-й и 3-й будут выбраны полностью.

n[1,…,1] – эквивалент n[1,:,:,1]

>>> n = numpy.arange(16).reshape(2, 2, 2, 2)
>>> n
array([[[[ 0,  1],
         [ 2,  3]],

        [[ 4,  5],
         [ 6,  7]]],


       [[[ 8,  9],
         [10, 11]],

        [[12, 13],
         [14, 15]]]])
>>> n[1,...,1]            # equivalent to n[1,:,:,1]
array([[ 9, 11],
       [13, 15]])
>>> # also Ellipsis object can be used interchangeably
>>> n[1, Ellipsis, 1]
array([[ 9, 11],
       [13, 15]])

Не больше, чем синтаксический сахар, да и тот со привкусом стевии.

Еще одно редкое применение (с версии питона 3.5) – многоточие может пригодиться в подсказках для типов.

а) Для указание типа кортежа неопределенной длины с однородными типами элементов:

Tuple[int, ...]

б) Для указания типа вызываемого объекта, когда сигнатура параметров неизвестна. Пример:

def partial(func: Callable[..., str], *args) -> Callable[..., str]:
    # Body

Это, конечно, уже экзотика.

Может быть у тебя, мой милый читатель, есть идеи, как еще применить…

 

 

Меняй значения переменных правильно

Достаточно одной строчки, временная переменная не нужна.

Специально для канала PyWay.

Defaultdict

Возьмем обычный питоновский dict. Определяем его мы так:

x = dict()

# или лучше

x = {}

Что будет, если мы обратимся к несуществующему элементу?

>>> x['abc']
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
KeyError: 'abc'

Возникает исключение KeyError. Можно попробовать отловить его конструкцией try-except, что выглядит и работает достаточно громоздко, можно проверить вхождение ключа операцией in или, наконец, воспользоваться методом get (в котором, кстати, мы можем указать значение по умолчанию, если ключа нет).

# плохо...
try:
    v = x['abc']
except:
    v = 'no value'

# лучше
v = x['abc'] if 'abc' in x else 'no value'

# еще лучше
v = x.get('abc', 'no value')

Это хорошо, но можно еще лучше!

Рассмотрим пример: допустим у нас есть задача по каждому ключевому слову из текста составить список позиций этого ключевого слова (номер слова) в тексте.

text = """
We develop a methodology for automatically analyzing text to aid in discriminating firms that encounter catastrophic 
financial events. The dictionaries we create from Management Discussion and Analysis Sections (MD&A) of 10-Ks 
discriminate fraudulent from non-fraudulent firms 75% of the time and bankrupt from nonbankrupt firms 80% of the 
time. Our results compare favorably with quantitative prediction methods. We further test for complementarities by 
merging quantitative data with text data. We achieve our best prediction results for both bankruptcy (83.87%) and 
fraud (81.97%) with the combined data, showing that that the text of the MD&A complements the quantitative financial 
information.
"""

key_words = [
    "quantitative",
    "results",
    "automatically"
]

# решение не претендует на общую эффективность, сделано для демонстрации

def solution1(text, keywords):
    # разбивка текста на слова и удаление лишних символов
    all_words = map(lambda word: word.strip(' .)(%\n').lower(), text.split(' '))

    kw_map = {}

    for word_no, word in enumerate(all_words):
        if word in key_words:
            if word in kw_map:
                kw_map[word].append(word_no)
            else:
                kw_map[word] = [word_no]

    return kw_map

print(solution1(text, key_words))

Нам понадобилось 4 строки, чтобы понять, есть ли уже такое слово словаре, если нет, то создать новый ключ со значением из списка с одним элементом, иначе добавить позицию в существующий список. Конечно, можно было сразу создать словарь, где ключи – ключевые слова, а значения – пустые списки, но есть и более элегантный способ – использовать defaultdict из модуля collection из стандартной библиотеки. Не зря мы в программировании всегда стремимся к простоте.

from collections import defaultdict

...

    kw_map = defaultdict(list)

    for word_no, word in enumerate(all_words):
        if word in key_words:
            kw_map[word].append(word_no)

    return kw_map

Ссылка на код.

Чем же отличается defaultdict от dict? И что значит параметр list?

Когда мы обращаемся к несуществующему ключу словаря, то defaultdict вызывает функцию, указанную при создании (в данном случае list) без параметров. И результат работы этой функции и будет присвоен новом элементу словаря с ключом, который раннее не существовал. А как только он появился, то исключении уже не будет брошено, и выполнится та операция, которую мы хотим (в данном случае append).

list – это встроенная функция, которая позволяет нам сконструировать список, если вызвана без параметров.

>>> print([])
[]
>>> print(list())   # тоже самое
[]

Поэкспериментируем в интерпретаторе:

>>> from collections import defaultdict
>>> d = defaultdict(list)   # создали нашdefaultdict
>>> d
defaultdict(<type 'list'>, {})
>>> d['test']   # усп! этого ключа еще нет в словаре
[]
>>> d   # а теперь есть, хотя мы всего лишь обратились на чтение
defaultdict(<type 'list'>, {'test': []})
>>> d['test2'].append('foo')   # test2 тоже нет, но он создасться как [] и append сработает
>>> d
defaultdict(<type 'list'>, {'test': [], 'test2': ['foo']})

Вместо list попробуем, например, int. То есть при обращении к несуществующему элементу defaultdict будет добавлено целое число (0 по умолчанию).

>>> d2 = defaultdict(int)
>>> d2['a']
0
>>> d2['b'] += 10   # операция += сработает, такое не пройдет с обычным dict
>>> d2
defaultdict(<type 'int'>, {'a': 0, 'b': 10})

defaultdict при создании принимает любую ф-цию, а не только встроенные! В следующем примеры мы создадим словарь по-умолчанию со значением элемента – числом 3 (с помощью лямбда ф-ции lambda: 3):

>>> d3 = defaultdict(lambda: 3)
>>> d3['fef']
3
>>> d3['rrr'] += 3
>>> d3
defaultdict(<function <lambda> at 0x10e0376e0>, {'fef': 3, 'rrr': 6})

И самое Питон-чудо напоследок! Смотрите как элегантно мы создадим тип данных «дерево» всего одной строкой!

def tree(): return defaultdict(tree)

Теперь мы можем:

x = tree()
x['a']['b']['c'] = 'test'
print(x)
defaultdict(<function tree at 0x10e037578>, {'a': defaultdict(<function tree at 0x10e037578>, {'b': defaultdict(<function tree at 0x10e037578>, {'c': 'test'})})})

Все поддеревья создались автоматически благодаря рекурсии! Каждый раз, когда мы хотим получить доступ к несуществующему поддереву, вызывается снова ф-ция tree, которая создает такое же поддерево, конструктор по-умолчанию которого тот же самый tree.

Специально для канала PyWay.