2012-06-04

リストの zip と unzip

関数プログラミングで基本的な関数の1つに zip がある。これはその名の通り、複数のリストをたばねて1つのリストにする関数である。このとき、入力リストの同じ位置の要素の組が出力リストの要素となる。例を Python で示すと次のようになる。

n = [1, 2, 3]
en = ['one', 'two', 'three']
fr = ['un', 'deux', 'trois']
z = zip(n, en, fr)
# z => [(1, 'one', 'un'), (2, 'two', 'deux'), (3, 'three', 'trois')]

今回の問題は、この zip されたリストが与えられたときに、元の複数のリストを返す関数 unzip はないんだろうか?ということである。Python には標準の組み込み関数として zip が用意されているが、unzip は見当たらない。どうしてだろう?と少し考えるとあることに気がついた。それは unzip は(ほとんど)zip だということだ。実際、次のようにすると元の複数のリストが得られる。

n, e, f = zip(*z)
# n => [1, 2, 3]
# e => ['one', 'two', 'three']
# f => ['un', 'deux', 'trois']

zip(*z) は Python 特有の記法で、リスト z の各要素を関数の引数として関数 zip を呼ぶことを表す。Lisp になじみがあるなら、次の記述の方がわかりやすいだろう。

n, e, f = apply(zip, z)

この unzip が実は zip だという事実は直感的にはやや不思議だが、次のように考えればあたりまえである。

[ 1    , 2     , 3       ],    [ [ 1, 'one'  , 'un'    ],
[ 'one', 'two' , 'three' ], =>   [ 2, 'two'  , 'deux'  ],
[ 'un' , 'deux', 'trois' ]       [ 3, 'three', 'trois' ] ]

見てすぐわかる通り、これは(ほぼ)行列の転置である。転置であるなら、もう1度適用すれば元に戻るのは自明である。

問題になるのは「ほぼ」と書いたところで、左の複数のリストが関数の複数の引数として表現されているのに対して右がリストのリストで複数のリストを表している。この zip の仕様のせいで、unzip するときに * (または apply) 相当の操作が必要となる。zip をリストのリストからリストのリストへの関数にしてしまえば、zip と unzip は完全に一致する。

無限リストの zip と unzip

リストが無限リストのときも、行列の転置だと考えれば zip = unzip ということができるのは自明だが、それを実装で確認することはできるだろうか?

例として自然数列と偶数列の zip を考えてみよう。Python には無限リストは無いので代わりにジェネレータを使うことにする。

以下のコードでは itertools を用い、以下の import 文を仮定する。

from itertools import count, tee, islice

まず無限リストの無限リストを表示で確認できるようにする。以下のコードでは、無限行列の左上 5 × 5 を表示している。

n = ( n for n in count() )
nn = ( n * 2 for n in count() )
m = ( y for y in [ n, nn ] )
print [ [ y for y in islice(x, 5) ] for x in islice(m, 5) ]
# => [[0, 1, 2, 3, 4], [0, 2, 4, 6, 8]]

zip を実装してみよう。ここでの zip は行を表すジェネレータのジェネレータ、つまり無限行列を引数にとり、それの転置行列を行を表すジェネレータのジェネレータとして返す。

def myzip(m):
    while True:
        mm, m = tee(m)
        yield ( y.next() for y in mm )

n = ( n for n in count() )
nn = ( n * 2 for n in count() )
m = ( y for y in [ n, nn ] )
t = myzip(m)
print [ [ y for y in islice(x, 5) ] for x in islice(t, 5) ]
# => [[0, 0], [1, 2], [2, 4], [3, 6], [4, 8]]

上のように、myzip で転置すると zip されていることがわかる。次に myzip を2回ほどこして元に戻るかどうかを試してみよう。

n = ( n for n in count() )
nn = ( n * 2 for n in count() )
m = ( y for y in [ n, nn ] )
t = myzip(myzip(m))
print [ [ y for y in islice(x, 5) ] for x in islice(m, 5) ]
# => [[0, 1, 2, 3, 4], [0, 2, 4, 6, 8], [], [], []]

無事に元に戻ることが確認された。

このように無限リストにおいても zip = unzip という考え方は有効である。

(注意) ジェネレータは一度消費すると二度と復元できないので、要所で tee を用いて複製する必要がある。myzip で tee しているところでは、本当はジェネレータが返すジェネレータも複製しなくてはならない。これは簡単にはできないような気がしたのであきらめた。このため、以下のようなことをすると誤動作する。

n = ( n for n in count() )
nn = ( n * 2 for n in count() )
m = ( y for y in [ n, nn ] )
t = myzip(m)
tt, t = tee(t)
print [ [ y for y in islice(x, 5) ] for x in islice(tt, 5) ]
m = myzip(t)
print [ [ y for y in islice(x, 5) ] for x in islice(m, 5) ]
# => [[], [], [], [], []]