【Python】zip, zip_longestの違い、同じ長さの入力を前提としたzip

本記事ではPythonにおいて複数の入力を列挙する関数であるzip、zip_longestおよびそれらの違いを紹介します。また、これらの関数は入力の長さが異なっていても動作するため、同じ長さを保証するように入力の要素を列挙する方法も紹介します。

zipの使用例

zipは与えられたリストや文字列などのイテラブルから要素のタプルを返すイテレータを作成します。

In [1]: for (x, y) in zip('abc', '123'):
   ...:     print(x, y)
   ...:
a 1
b 2
c 3

zipはイテラブルの中で最短の長さのイテレータを返す点に注意が必要です。以下の様にabc、12の二つの文字列を与えると、長さが2のイテレータを作成します。

In [2]: for (x, y) in zip('abc', '12'):
   ...:     print(x, y)
   ...:
a 1
b 2

Pythonのドキュメントではzipに関して以下のように説明があります。

zip() should only be used with unequal length inputs when you don’t care about trailing, unmatched values from the longer iterables. If those values are important, use itertools.zip_longest() instead.

長さが異なる入力に対してzipを使う場合は、長いイテラブルの終わりのほうの要素が使われないことに注意が必要です。それが困る場合はitertools.zip_longestを使うことが推奨されています。

itertools.zip_longestの使用例

zipと似たzip_longestという関数も利用できます。 zip_longestはzipと異なり、最長のイテラブルの長さのイテレータを作成します。最長のイテラブル以外のものについて、長さを超えた場合の要素はデフォルトではNoneが返ってきます。

In [3]: import itertools

In [4]: for (x, y) in itertools.zip_longest('abc', '12'):
   ...:     print(x, y)
   ...:
a 1
b 2
c None

None以外の他の値を返すようにしたい場合はfillvalueという引数を与えます。

In [5]: for (x, y) in itertools.zip_longest('abc', '12', fillvalue='O'):
   ...:     print(x, y)
   ...:
a 1
b 2
c O

上記のようにitertools.zip_longestでは最長の入力に合わせてイテレータを作成することができます。

iterablesが異なる長さの場合エラーを出すようにzip_longestを使用する

上記で紹介したように、zipは入力の長さが異なっていた場合に、最短の入力の長さのイテレータを作成します。一方でitertools.zip_longestは最長の入力の長さのイテレータを作成します。そのため、これら2つの関数を利用するときに、入力の長さが同じことを想定しているにもかかわらず、入力の長さが異なっていても、プログラムは動作する可能性があり、バグの原因となることもあると思います。

そこで、zip_longestに与える入力の長さが等しいことを保証したい場合を考えます。単純にはzip_longestに入力する前にすべてのイテラブルの長さが等しいことを確認すれば良いですが、イテレータを入力とする場合、len()で長さを取得できません。

イテレータが入力として与えられても、長さが等しくない場合にエラーを出せるようにするために、zip_longestを使って以下の様な関数を作成します。

In [1]: import itertools

In [2]: def zip_equal(*iterables):
   ...:     sentinel = object()
   ...:     for combination in itertools.zip_longest(*iterables, fillvalue=sentinel):
   ...:         if sentinel in combination:
   ...:             raise ValueError('iterables have different lengths')
   ...:         yield combination
   ...:

この関数ではfillvalueで指定した要素がzip_longestによって得られたタプルに含まれていた場合、エラーを出すようにしています。fillvalueがあるということはイテラブルの長さが等しくないことを表すため、このエラーはイテラブルの長さが等しくない場合に出ます。

長さが等しい入力とそうでない入力を与えて結果を確認します。

In [3]: for (x, y) in zip_equal('abc', '123'):
   ...:     print(x, y)
   ...:
a 1
b 2
c 3

In [4]: for (x, y) in zip_equal('abc', '12'):
   ...:     print(x, y)
   ...:
a 1
b 2
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-4-d2760dd8dd39> in <module>
----> 1 for (x, y) in zip_equal('abc', '12'):
2     print(x, y)
3

<ipython-input-2-2f6da01ba3db> in zip_equal(*iterables)
3     for combination in itertools.zip_longest(*iterables, fillvalue=sentinel):
4         if sentinel in combination:
----> 5             raise ValueError('iterables have different lengths')
6         yield combination
7

ValueError: iterables have different lengths