イテレータとジェネレータ¶

In [1]:

import sys

イテレータ¶

自然数$n$が与えられたとき、${}_n C_k$ ($k = 0, 1, \dots, n$)、および$\sum_{k=0}^{n} {}_n C_k$を計算するプログラムの例を示す。

In [2]:

n = 6

In [3]:

v = 1
s = 0
for k in range(n+1):
    print(v)
    s += v
    v *= (n-k)
    v //= (k+1)
s

Out[3]:

このプログラムは期待した通りの結果を返すが、for文の反復内容に${}_n C_k$を計算する部分が入るため、見通しが悪い。また、$\sum_{k=0}^{n} (-1)^k{}_n C_k$を計算するようにプログラムを改変するには、${}_n C_k$を計算する部分も含めて、実装を修正する必要が生じる。見通しのよいプログラムを書くために、range関数が$0, 1, \dots, n$を返すように、${}_n C_0, {}_n C_1, \dots, {}_n C_n$を返す関数combinationsを定義し、以下のようなプログラムを書きたい。

s = 0
for v in combinations(n):
    print(v)
    s += v
s

おそらく、真っ先に思い浮かぶのはcombinations関数を${}_n C_k$ ($k = 0, 1, \dots, n$)のリストを返す関数として実装することであろう（後ほど異なる実装を示すことになるので、ここではcombinationではなくlist_combinationsという名前の関数として実装する）。

In [4]:

def list_combinations(n):
    v = 1
    C = []
    for k in range(n+1):
        C.append(v)
        v *= (n-k)
        v //= (k+1)
    return C

list_combinations関数の動作を確認する。

In [5]:

list_combinations(n)

Out[5]:

[1, 6, 15, 20, 15, 6, 1]

このlist_combinations関数を使うと、先ほどの処理を行うプログラムは次のように書ける（呼び出し元のプログラムの見通しは良くなった）。

In [6]:

s = 0
for v in list_combinations(n):
    print(v)
    s += v
s

Out[6]:

また、$\sum_{k=0}^{n} (-1)^k{}_n C_k$を計算するには、その計算の箇所だけ変更すればよい（動作を分かりやすくするため、print(k, v)を入れてある）。

In [7]:

t = 0
for k, v in enumerate(list_combinations(n)):
    print(k, v)
    t += ((-1) ** k * v)
t

Out[7]:

なお、和の計算結果のみが必要であれば、sum関数を用いるだけでよい。

In [8]:

sum(list_combinations(n))

Out[8]:

これにより、見通しのよいプログラムを分かりやすく書くことができたが、一つだけ心残りがある。それは、${}_n C_0, {}_n C_1, \dots, {}_n C_n$のすべての値を要素とするリストを一時的に作成・保持してしまう点である。今回の目的が$\sum_{k=0}^{n} {}_n C_k$の値を計算することだけであったら、${}_n C_k$のすべての値をメモリに保持するのは無駄である。なぜなら、${}_n C_k$のすべての値に同時またはランダムにアクセスする必要は無いからである。

${}_n C_0, {}_n C_1, \dots, {}_n C_n$のように、要素の値を順に返すような動作を実現してくれる仕組みがイテレータである。オブジェクトにイテレータの機能を実装することで、for文の反復対象として用いることができるようになる。オブジェクトにイテレータを実装するには、__iter__メソッドと__next__メソッドを実装すればよい。__iter__メソッドはイテレータオブジェクト自身を返すもので、通常はselfを返せばよい。__next__メソッドには次の要素を返す処理を実装する。これ以上返す要素が無い場合は、StopIteration例外を送出させる。

In [9]:

class iter_combinations:
    def __init__(self, n):
        self.v = 1
        self.k = 0
        self.n = n
    
    def __iter__(self):
        return self
    
    def __next__(self):
        if self.n < self.k:
            raise StopIteration
        v = self.v
        self.v *= (self.n - self.k)
        self.v //= (self.k+1)
        self.k += 1
        return v

このiter_combinations関数（実際にはクラス）を呼び出すプログラムは、list_combinations関数でリストを予め生成した場合と完全に同じである。

In [10]:

s = 0
for v in iter_combinations(n):
    print(v)
    s += v
s

Out[10]:

なお、和の計算結果のみが必要であれば、list_combinations関数でリストを予め生成した場合と同様に、sum関数を用いるだけでよい。

In [11]:

sum(iter_combinations(n))

Out[11]:

$\sum_{k=0}^{n} (-1)^k{}_n C_k$の計算も同様である。

In [12]:

t = 0
for k, v in enumerate(iter_combinations(n)):
    print(k, v)
    t += ((-1) ** k * v)
t

Out[12]:

当たり前ではあるが、iter_combinationsクラスのオブジェクトはリストではない。

In [13]:

iter_combinations(n)

Out[13]:

<__main__.iter_combinations at 0x7f07a80aae30>

In [14]:

type(iter_combinations(n))

Out[14]:

__main__.iter_combinations

In [15]:

isinstance(iter_combinations(n), list)

Out[15]:

False

iter_combinationsクラスのイテレータ・オブジェクトから順に要素を取り出すにはnext関数を用いる。すべての要素を取り出した後、StopIteration例外が送出されていることが分かる。

In [16]:

i = iter_combinations(4)

In [17]:

next(i)

Out[17]:

In [18]:

next(i)

Out[18]:

In [19]:

next(i)

Out[19]:

In [20]:

next(i)

Out[20]:

In [21]:

next(i)

Out[21]:

In [22]:

next(i)

---------------------------------------------------------------------------
StopIteration                             Traceback (most recent call last)
/tmp/ipykernel_205148/1939748483.py in <module>
----> 1 next(i)

/tmp/ipykernel_205148/1028479797.py in __next__(self)
     10     def __next__(self):
     11         if self.n < self.k:
---> 12             raise StopIteration
     13         v = self.v
     14         self.v *= (self.n - self.k)

StopIteration:

実際にイテレータ・オブジェクトが消費しているメモリ量を見積もるため、sys.getsizeofを用いる。$n=5$のとき、list_combinations関数が生成するオブジェクトよりもiter_combinationsクラスが生成するオブジェクトの方がメモリ消費量が少ない。

In [23]:

sys.getsizeof(iter_combinations(5))

Out[23]:

In [24]:

sys.getsizeof(list_combinations(5))

Out[24]:

続いて、$n=100000$とすると、メモリ消費量の差は顕著になる。list_combinations関数が生成するオブジェクトのメモリ消費量は増えていくが、iter_combinationsクラスが生成するオブジェクトのメモリ消費量は$n$によらず一定である。これは、イテレータがすべての要素を持つリストを生成せずに、要素を順番に返す仕様になっているからである。

In [25]:

sys.getsizeof(list_combinations(100000))

Out[25]:

In [26]:

sys.getsizeof(iter_combinations(100000))

Out[26]:

ジェネレータ¶

先ほどの${}_n C_k$ ($k = 0, 1, \dots, n$)を計算する例において、list_combinations関数とiter_combinations関数の実装を再掲する。。

In [27]:

def list_combinations(n):
    v = 1
    C = []
    for k in range(n+1):
        C.append(v)
        v *= (n-k)
        v //= (k+1)
    return C

class iter_combinations:
    def __init__(self, n):
        self.v = 1
        self.k = 0
        self.n = n
    
    def __iter__(self):
        return self
    
    def __next__(self):
        if self.n < self.k:
            raise StopIteration
        v = self.v
        self.v *= (self.n - self.k)
        self.v //= (self.k+1)
        self.k += 1
        return v

s = 0
#for v in list_combinations(6):
for v in iter_combinations(6):
    print(v)
    s += v
s

Out[27]:

list_combinations関数とiter_combinationsクラスの実装を比較すると、list_combinations関数の方が見通しがよい。これは、反復処理をfor文で書くことができるため、反復の開始から終了まで、繰り返しの処理内容をまとめて記述できるからである。一方、iter_combinationsクラスにおける__next__メソッドでは、現在の状態から次の要素を返す処理を記述することになるため、反復処理全体の流れが分かりにくい。

この両者の良いところ取りはできないものか？　すなわち、list_combinations関数のように反復処理を書くことができ、かつiter_combinationsクラスのようにイテレータとして振る舞う方法は無いだろうか？　これを実現するのが、yield文によるジェネレータ関数である。関数の中でyield文を記述すると、その関数はジェネレータ関数となる。yield文はreturn文に似ているが、関数の振る舞いをイテレータに変更し、要素を順番に返せるようにするものである。yield文が実行されると、関数の実行は中断し、yield文で指定された値が要素として呼び出し元に返される。言い換えれば、yield文の実行と、イテレータの__next__メソッドの返り値が対応する。

以下のgen_combinations関数は、ジェネレータ関数によるイテレータの実装である。list_combinations関数でリストにappendする代わりに、yield文を実行するプログラムとなる。

In [28]:

def gen_combinations(n):
    v = 1
    for k in range(n+1):
        yield v
        v *= (n-k)
        v //= (k+1)

s = 0
for v in gen_combinations(6):
    print(v)
    s += v
s

Out[28]:

In [29]:

sum(gen_combinations(6))

Out[29]:

yield文を含む関数はジェネレータ関数となる。関数の戻り値はジェネレータ型であり、これはイテレータ・オブジェクトとして用いることができる。

In [30]:

gen_combinations(4)

Out[30]:

<generator object gen_combinations at 0x7f07a816a0a0>

In [31]:

type(gen_combinations(4))

Out[31]:

generator

ジェネレータ関数の返り値はイテレータとして振る舞うので、next関数を用いて順に要素を取り出すことができる。すべての要素を取り出した後、StopIteration例外が送出される。

In [32]:

i = gen_combinations(4)

In [33]:

next(i)

Out[33]:

In [34]:

next(i)

Out[34]:

In [35]:

next(i)

Out[35]:

In [36]:

next(i)

Out[36]:

In [37]:

next(i)

Out[37]:

In [38]:

next(i)

---------------------------------------------------------------------------
StopIteration                             Traceback (most recent call last)
/tmp/ipykernel_205148/1939748483.py in <module>
----> 1 next(i)

StopIteration:

ジェネレータ式¶

リストの内包表記で$\{n^2 | n \in \{0, 1, 2, 3, 4, 5\}\}$をリストとして得る例である。

In [39]:

[n**2 for n in range(6)]

Out[39]:

[0, 1, 4, 9, 16, 25]

得られたリストに対して、sum関数で和を取ることで、$\sum_{n=0}^5 n^2$を計算する。

In [40]:

sum([n**2 for n in range(6)])

Out[40]:

[]を()に変更すると、ジェネレータを内包表記で作成できる（ジェネレータ式）。以下のように関数の唯一の引数として用いる場合は、括弧を二重（((や))）にしなくてもよい。

In [41]:

sum(i**2 for i in range(6))

Out[41]:

ジェネレータ式で$\sum_{k=0}^{6} (-1)^k{}_6 C_k$を計算する。

In [42]:

sum((-1) ** k * v for k, v in enumerate(iter_combinations(6)))

Out[42]: