【python】FiniteSetを使って組み合わせを一気に作る

Pythonから始める数学入門」の演習課題で、「トランプのカードをきる」というのがある。 カードは絵柄と数字の組み合わせを持っていて、これを52枚用意してランダムにシャッフルした上で結果を表示するというもの。

52枚を人力で丹精込めて1枚ずつリストに込めてもできる。 が、集合演算を使って組み合わせをパパッとつくってループを回すことで簡単にできた。 ちょっと感動したので共有。

from sympy import FiniteSet
import random

class Card:
    def __init__(self, suit, rank):
        self.suit = suit
        self.rank = rank

suits = FiniteSet('diamonds','hearts','clubs','spades')
ranks = FiniteSet('ace',2,3,4,5,6,7,8,9,10,'jack','queen','king')
cards = []
for c in suits*ranks:
    cards.append(Card(c[0],c[1]))
random.shuffle(cards)
for c in cards:
    print('{0} of {1}'.format(c.rank, c.suit))

今回感動したのは、絵柄と数字の組み合わせを直積で計算してループを回しているところ。 FiniteSet * FiniteSet とすることで直積が計算でき、これによって2つの集合からそれぞれ要素を選んだ場合のすべての組み合わせを求めることができる。

直積については、こんな感じになる。

>>> from sympy import FiniteSet
>>> s = FiniteSet(1,2)
>>> t = FiniteSet(3,4)
>>> s * t
{1, 2} x {3, 4}
>>> for u in s*t:
...     print(u)
...
(1, 3)
(1, 4)
(2, 3)
(2, 4)

これを使えば2つのリストから各要素を取り出して組み合わせを作ってそれを元に計算して…なんてことを結構簡単にできる。