【AtCoder版！蟻本】JOI 2009 予選D カード並べ【順列全探索】

AtCoder版！蟻本のn!通り全探索類題としてあげられているJOI 2009 予選D カード並べをPythonで解いていきます。

リンク

D - カード並べ

D - カード並べ

問題文

花子さんは n 枚 (4≤n≤10) のカードを並べて 遊んでいる． それぞれのカードには 1 以上 99 以下の整数 が 1 つずつ書かれている． 花子さんは，これらのカードの中から k 枚 (2≤k≤4) を選び， 横一列に並べて整数を作ることにした． 花子さんは全部で何種類の整数を作ることができるだろうか．
例えば， 1,2,3,13,21 の 5 枚のカードが与えられ，そこから 3 枚を選び整数を作ることを考える．
2,1,13 をこの順に並べると，整数 2113 を作ることができる． また， 21,1,3 をこの順に並べても，同じ整数 2113 を作ることができる． このように，異なるカードの組み合わせから同じ整数が作られることもある．
n 枚のカードに書かれた整数が与えられたとき， その中から k 枚を選び，横一列に並べることで作ることができる整数の個数を求めるプログ ラムを作成せよ．

制約

・入力は 2+n 行からなる．
1 行目にはカードの枚数 n (4≤n≤10) が， 2 行目にはカードを選ぶ 枚数 k (2≤k≤4) が書かれている．
2+i 行目 (1≤i≤n) には i 枚目のカードに書かれて いる 1 以上 99 以下の整数が書かれている．

解答例

import itertools

n = int(input())
k = int(input())

num = []
for i in range(n):
    num.append(int(input()))

ans = set()

for i in itertools.permutations(range(n)):
    tmp = 0
    for j in range(k):
        if num[i[j]] < 10:
            tmp = tmp*10 + num[i[j]]
        elif num[i[j]] >= 10:
            tmp = tmp*100 + num[i[j]]
    ans.add(tmp)

print(len(ans))

解説

1から99の整数が書かれたカードがn枚あり、そのなかからk枚を並べて整数を作る場合に全部で何種類の整数を作ることができるか答える問題です。

順列を用いてカードの並び方を作成し、そこから作ることのできる整数を全探索しました。

解答になるansは重複しない要素を集めたいのでset型です。

iがn枚のカードの順列になります。
このすべてのパターンで最初からk枚のカードを選ぶことで整数を作っていきました。

tmp変数に作る整数を格納していきます。
選んだカードの数字が10より小さければtmpに10をかけて数字を足します。
10以上ならtmpに100をかけて足すことでtmpを更新していきました。

k枚のカードを選び終えたあとにansにtmpを加えていきます。

ループを抜けた後にansの長さを解答すればACです！

AtCoder版！蟻本（初級編）の記事リンクページはこちら
ebisuke33.hatenablog.com

リンク