Python3で解く AtCoder Beginner Contest 184 D - increment of coins

単純な確率問題でラッキーと思ったら、しっかりDPとの組み合わせだった。（そりゃそうだDだもの）

目次

概要

問題

• 金貨が枚、銀貨が枚、銅貨が枚袋に入ってる。
• 袋の中にあるいずれかの種類の硬貨が 枚になるまで以下の操作を繰り返す。
  • ランダムに袋から枚取る。
  • 取ったものと同じ種類の硬貨を枚袋に戻す。（要するに取ったものが枚増える）
    （この意味が最初わからかなかった。どこから硬貨が湧いて出てくるのかとかを不必要に考えてしまった。）
• 操作回数の期待値は？

---

解くときに考えた内容

• 金貨を引く確率×金貨を引いた時の操作回数の期待値＋ 銀貨を引く確率×銀貨を引いた時の操作回数の期待値＋ 銅貨を引く確率×銅貨を引いた時の操作回数の期待値というのをDPで解く。
• 上記は、再帰で解く必要がある。関数の中身は、
  • いずれかが100枚だったら、0を返すようにしておく。
    • 99枚のときに引いた時の操作回数の期待値が0になるイメージ。
    • すべての硬貨が99枚のときはどれを引いてもあと1回で終了。
    • つまり、再帰の一番の元はmemo[99][99][99] = 1ってこと。
• あとは、メモ化をきちんとしておけばよい。
• 普段DPは0, 1, 2などを初期値を規定して所定の値の結果を求めるイメージだが、今回は再帰ですべてが99の状態を起点に元に逆流して答えが求まるイメージで組み立てる。
• 今回この問題を解くにあたり、functools.lru_cacheという便利な関数があることを知った。
  docs.python.org 再帰の関数を@lru_cacheとデコレータするだけで、自動的にメモ化してくれる。実装例を見るとmemoのところをすべてすっ飛ばして良くなるので超楽ちん。

---

コード

リストでメモ化

テストのhand_03.txtの結果。

def memo_exp(a, b, c):
    exp = memo[a][b][c]
    if exp != None:
        return exp
    if a == 100 or b == 100 or c == 100:
        return 0
    sum_all = a + b + c
    memo[a][b][c] = 1 + a / sum_all * memo_exp(a + 1, b, c) + \
        b / sum_all * memo_exp(a, b + 1, c) + \
        c / sum_all * memo_exp(a, b, c + 1)
    return memo[a][b][c]

a, b, c = map(int, input().split())
memo = [[[None] * 101 for _ in range(101)] for _ in range(101)]
print(memo_exp(a, b, c))

テストのhand_03.txtの結果。

lru_cache（超便利）

テストのhand_03.txtの結果。 memo化と比べると、時間もメモリも食うのがわかる。
（今回は1例だけ上げていているがほとんどのサンプルで同じ傾向） なので、メモリと時間的制約の余裕があればこちらで良さそう。

from functools import lru_cache
@lru_cache(None)
def lru_exp(a, b, c):
    if a == 100 or b == 100 or c == 100:
        return 0
    sum_all = a + b + c
    return 1 + a / sum_all * lru_exp(a + 1, b, c) + \
        b / sum_all * lru_exp(a, b + 1, c) + \
        c / sum_all * lru_exp(a, b, c + 1)
a, b, c = map(int, input().split())
print(lru_exp(a, b, c))

---

参考になる書籍

• 問題解決力を鍛える!アルゴリズムとデータ構造

けんちょんさんこと、大槻さんの書籍である、「問題解決力を鍛える!アルゴリズムとデータ構造」 ご本人の記事の中で難易度は、螺旋本ということで初学者でも非常に結構読みやすい!!!（実際に、けんちょんさんの記事はC++に疎い自分でも、その思考プロセスを学べるくらい丁寧に書かれているし、この書籍も非常に読みやすい）どのようなステータスの人にもおすすめの一冊になっている。

とうとう発売され自分の手元にも来た（2020-10-03）。現在読み中なのである程度読み込んだら感想を書く予定。めっちゃ図が多いし、章末問題もあるしかなりイケてる。答えはC++で書かれているがそのうちPythonでも書かれるとのことなので、自分で前もってやってしまおうと思う（答え合わせできるのはありがたい）。

問題解決力を鍛える!アルゴリズムとデータ構造 (KS情報科学専門書)

７章までは別記事に残してある。 blacktanktop.hatenablog.com