Python標準ライブラリを使いこなす 初級編 ~二重でfor文書いたりしてませんか?~

概要

Pythonは標準ライブラリで便利な機能を多数用意しています。

しかし、きちんと使いこなせている人は少ないのではないでしょうか。可読性の高いコードを書くことができるように標準ライブラリを使った方法に書き直す問題を作りました。

問題

二重ループ

悪い例

for i in range(100):
    for j in range(1000):
        # some process
        pass

二重ループを上の例のように書くと実際の処理の部分のネストが深くなってしまいます。Pythonはネストを4文字分のスペースで表現するので、ネストが深くなってしまうとline lengthが非常に長くなってしまい可読性が低下してしまいます。
独立な変数のループの場合、標準ライブラリを使って一段階のネストでこのループを表現する方法があります。
上記の悪い例を書き直してみましょう。

回答

from itertools import product

for i, j in product(range(100), range(1000)):
    # some process
    pass

itertoolsのproductを利用することで1個のネストで多重ループを表現できます。
productは二重ループだけでなく同様の方法で$n$重ループに使用できます。

ただし、この方法が利用できるのはそれぞれの変数が独立な場合のみです。

また、同じイテラブルの多重ループの場合、以下のように書くこともできます。

from itertools import product

for i, j in product(range(100), repeat=2):
    # some process
    pass

ループで前後のアイテムを利用

悪い例

iterator = iter(some_iterable)
old = next(iterator)
for new in iterator:
    # some process
    old = new

イテラブルの前後を利用するループは比較的登場回数が多いと思います。
上の例ではoldやnewの値を書き換えたりcontinue文を利用すると後のループに影響を与えてしまう問題があります。
上記の悪い例を書き直してみましょう。

回答

from itertools import pairwise

for old, new in pairwise(some_iterable):
    # some process
    pass

itertoolsのpairwiseを利用することで連続する要素をペアで取り出すことができます。

0から無限までの整数の生成

悪い例

i = 0
while True:
    # some process
    i += 1

0から無限までの整数の生成してbreak文と組み合わせたいことがあると思います。
上の例ではwhile文を利用して、0から無限までの整数の生成しています。
上の例も前後のアイテムを利用する例と同様に、iを書き換えたりcontinue文を利用すると後のループに影響を与えてしまう問題があります。
上記の悪い例を書き直してみましょう。

回答

from itertools import count

for i in count():
    # some process
    pass

itertoolsのcountを利用して、0から無限までの整数の生成することができます。

上位10件を取得

和が100になる3つの数$i, j, k (i\ge j \ge k)$の組み合わせのうち積が大きいものそれぞれ上位10件取得してみましょう。

悪い例

from math import prod

n = 100
candidates = [
    (i, j, k) for i in range(1, n - 1) for j in range(1, n - i) if i >= j >= (k := n - i - j)
]
result = sorted(candidates, key=prod, reverse=True)[:10]

Pythonの比較演算子は3個以上の変数を比較するときi >= j >= kのような記法が利用できます。
:=は代入式です。代入を文ではなく式で行うことができます。
上位複数件を取得したい場合、sortedとスライスの組み合わせは上位以外のアイテムも並べ替えるため余分な計算をしていることになってしまいます。
その部分を書き直してみましょう。

回答

from heapq import nlargest
from math import prod

n = 100
candidates = [
    (i, j, k) for i in range(1, n - 1) for j in range(1, n - i) if i >= j >= (k := n - i - j)
]
result = nlargest(10, candidates, key=prod)

heapqのnlargestを利用して、上位複数件を取得することができます。

積集合

悪い例

sets = [{1, 2, 3}, {2, 8}, {2, 4, 6}, {1, 2}]
iterator = iter(sets)
prod_set = next(iterator)
for num_set in iterator:
    prod_set &= num_set

上記のような累積の演算を行うものをきれいにまとめる方法が標準ライブラリには用意されています。
上記の悪い例を書き直してみましょう。

回答

from functools import reduce
from operator import and_

sets = [{1, 2, 3}, {2}, {2, 4, 6}, {1, 2}]
prod_set = reduce(and_, sets)

functoolsのreduceとoperatorのand_を利用します。
組み込み関数のsumやmathのprodのような累積の演算を任意の二項演算に行うユースケースではfunctoolsのreduceを利用します。頻出なので覚えておいて損はないです。
operatorはpythonの演算子を関数形式で提供している標準ライブラリです。例えばand_はlambda a, b: a & bと同様の役割ですがより高速に処理できます。

最後に

今回の問題では使用頻度が高い簡単なものだけを紹介しました。
また、もっときれいに書く方法があれば教えてもらえると嬉しいです。