ビンゴの確率を計算する

bingo_prob.py

import math
from fractions import Fraction

bingo5x5_seqs = [
  [1, 2, 3, 4, 5],
  [6, 7, 8, 9, 10],
  [11, 12, 13, 14],
  [15, 16, 17, 18, 19],
  [20, 21, 22, 23, 24],
  [1, 6, 11, 15, 20],
  [2, 7, 12, 16, 21],
  [3, 8, 17, 22],
  [4, 9, 13, 18, 23],
  [5, 10, 14, 19, 24],
  [1, 7, 18, 24],
  [5, 9, 16, 20],
]

def is_bingo(pattern: list[bool], bingo_seqs: list[list[int]]) -> bool:
  """
  ビンゴかどうか判定する.

  Args:
    pattern: ビンゴカードのマスの穴空きパターン. Trueならマスが空いている.
    bingo_seqs: ビンゴのパターン. listに含まれる全ての番号のマスが空いていれば、ビンゴと判定する.
  Returns:
    ビンゴならTrue
  """
  return any(all(pattern[i-1] for i in seq) for seq in bingo_seqs)


def build_card_pattern(n: int, mask: list[int]) -> list[bool]:
  """
  ビンゴカードの穴空きパターンを作る.

  Args:
    n: 全列挙中のパターンの番号. 0 <= n < 2^24

  Returns:
    真理値のリスト. Trueならマスが空いている
  """
  return [(n & m) != 0 for m in mask]


def count_patterns(n: int, bingo_seqs: list[list[int]]) -> list[int]:
  """
  ビンゴカードの空いているマスの数ごとに、ビンゴになっているパターンの数を数え上げる.

  Args:
    n: マス目の数 (24)
    bingo_seqs: ビンゴのパターン. listに含まれる全ての番号のマスが空いていれば、ビンゴと判定する.
  Returns:
    リスト. indexがiの要素は、i個のマスが空いているパターンの数を表す.
  """
  counts = [0] * (n + 1) # counts[i]: i個のマスが空いているパターンの数

  mask = [1 << i for i in range(n)]
  for i in range(1 << n):
    pattern = build_card_pattern(i, mask)

    if is_bingo(pattern, bingo_seqs):
      counts[i.bit_count()] += 1

  return counts

def calc_P_given_Hk(k: int, counts: list[int]) -> Fraction:
  """
  ビンゴカードのマスがk個空いている時に、ビンゴになっている確率(P[B|H_k])を計算する.

  Args:
    k: 空いているマスの数
    counts: 空いているマスの数ごとに、ビンゴになっているパターンの数を数え上げたリスト.
  """
  return Fraction(counts[k] * math.factorial(k), math.perm(24, k))


def calc_P_K_Hk(N: int, n: int, K: int, k: int) -> Fraction:
  """
  K回目の抽選後に、マスがk個空いている確率(P_K[H_k])を計算する.

  Args:
    N: 抽選総数 (75)
    n: マスの数 (24)
    K: 抽選回数
    k: 空いたマスの数
  """
  if K > N:
    raise Exception("抽選総数より抽選回数が多い")
  if k > n:
    raise Exception("マスの数より空いているマスの数が多い")
  if (K < k) or ((n - k) > (N - K)):
    return Fraction(0, 1)

  return Fraction(math.comb(K, k) * math.perm(n, k) * math.perm(N-n, K-k), math.perm(N, K))


def calc_P_K_B(N: int, n: int, K: int, ps: list[Fraction]) -> Fraction:
  """
  K回目の抽選後にビンゴになっている確率(P_K[B])を計算する.

  Args:
    N: 抽選総数 (75)
    n: マスの数 (24)
    K: 抽選回数
    ps: マスがk個空いている時、それがビンゴである確率(P[B|H_k])
  """
  ret = Fraction(0, 1)
  for k in range(min(K, n) + 1):
    p = calc_P_K_Hk(N, n, K, k) * ps[k]
    ret += p

  return ret


def calc_from_sum(ps: list[Fraction]) -> list[Fraction]:
  """
  累積された確率のリストから、元の確率のリストを計算する.

  Args:
    ps: 累積された確率のリスト
  """
  return [ps[0]] + [ps[i] - ps[i-1] for i in range(1, len(ps))]


def calc_E(ps: list[Fraction]) -> Fraction:
  """
  期待値を計算する
  """
  ret = Fraction(0, 1)
  for i in range(len(ps)):
    ret += i * ps[i]

  return ret


# k個のマスの空け方
counts = count_patterns(24, bingo5x5_seqs)
for i, c in enumerate(counts):
  print(i, c)

# k個マスが空いている時に、ビンゴになっている確率(P[B|H_k])
ps_given_Hk = [calc_P_given_Hk(i, counts) for i in range(len(counts))]
for i, p in enumerate(ps_given_Hk):
  print(i, p, p.numerator / p.denominator)

# 丁度k個目のマスを空けた時にビンゴになる確率(P[B|H_k]')
ps_given_Hk_ = calc_from_sum(ps_given_Hk)
for i, p in enumerate(ps_given_Hk_):
  print(i, p, p.numerator / p.denominator)

# E_k
print(calc_E(ps_given_Hk_))

# K回目の抽選後にビンゴになっている確率(P_K[B])
ps_B = [calc_P_K_B(75, 24, i, ps_given_Hk) for i in range(76)]
for i, p in enumerate(ps_B):
  print(i, p)

# 丁度K回目の抽選でビンゴになる確率(P_K[B]')
ps = calc_from_sum(ps_B)
for i, p in enumerate(ps):
  print(i, p)

# E_K
print(calc_E(ps))