AHC030をPythonでゆっくりやってみる

AHC030のベイズ推定を使う解法をPythonでやってみます。 以下の記事を参考にさせていただきました。 qiita.com

M=2のときの配置推定

まずはM=2のときの配置推定（ベイズの定理）からやってみます。

ベイズの定理の細かい内容は飛ばします。（僕もお気持ち程度しかわかっていません）

とりあえず、油田の配置可能集合とし、これまでマス集合に対して占いをし結果がであった時のが知りたいです。

つまり、真の配置がであるとき、占い結果がになる確率を計算します。

これを計算するには、「マスの集合を占い、真の石油埋蔵量がのときに占い結果がとなる確率」を考えればよいです。

したがって、真の配置がのときのマス集合の石油埋蔵量をと書くと、となります。

実装時には、占い結果が近似値で入力されることに注意が必要です。

例えば、占い結果が1.3のときジャッジプログラムからは1が返されます。

そのため、確率を考えるときは～の範囲を考える必要があります。

import math
from random import randint

N,M,E = map(float, input().split())
N,M = int(N),int(M)
D = list()
SUM_D = 0
IJ = list()
MIN_D = 10**18
for _ in range(M):
  d,*ij = map(int,input().split())
  MIN_D = min(MIN_D,d)
  D.append(d)
  SUM_D += d
  L = list()
  for ind in range(0,2*d,2):
    L.append((ij[ind],ij[ind+1]))
  IJ.append(L)
for ind in range(M): # 各油田が左上から順に並ぶようにする
  IJ[ind].sort()

def make_all_board(): # 配置可能な油田位置を列挙する（ここはもっといい書き方があると思います）
  BOARD = list()
  m_board = list() # [1番目の油田の配置可能位置,2番目の油田の配置可能位置]
  for m in range(M):
    tmp = list()
    for i in range(N):
      for j in range(N): # m番目の油田の左上を(i,j)に合わせる
        nn = set()
        ng = False
        for ii,jj in IJ[m]:
          ni,nj = i+ii,j+jj
          if not (0<=ni<N and 0<=nj<N): 
            ng = True
            break
          nn.add((ni,nj))
        if ng: continue
        tmp.append(nn)
    m_board.append(tmp)
  for m1 in range(len(m_board[0])):
    for m2 in range(len(m_board[1])):
      nn = [[0]*N for _ in range(N)]
      for i,j in m_board[0][m1]:
        nn[i][j] += 1
      for i,j in m_board[1][m2]:
        nn[i][j] += 1
      BOARD.append(nn)
  
  return BOARD

def normal_cdf(x,mean,sigma):
  return 0.5 * (1.0 + math.erf((x - mean) / (sigma * math.sqrt(2))))

def probability_in_range(l,r,mean,sigma):
  assert l <= r
  if mean < l:
    return probability_in_range(2.0 * mean - r, 2.0 * mean - l, mean, sigma)
  p_l = normal_cdf(l, mean, sigma)
  p_r = normal_cdf(r, mean, sigma)
  return p_r - p_l

def q(k,s,r): # kマスの集合Sを占い、石油埋蔵量がsのとき、占い結果がrとなる確率を返す
  # 問題文に従って平均,分散を計算
  mean = (k-s)*E + s*(1-E)
  sigma = (k*E*(1-E))**0.5
  
  # 占い結果は近似値なので、+-0.5の範囲を調べる
  if r==0: return probability_in_range(-1e10, r+0.5, mean, sigma) # 0の時だけ-∞にする必要がある
  else: return probability_in_range(r-0.5, r+0.5, mean, sigma)
  
def normalize(p): # 事後確率の計算
  SUM = sum(p)
  assert SUM>0
  return [x/SUM for x in p]

def bayesian_inference(BOARD):
  n = len(BOARD)
  P = [1.0/n]*n
  while True:
    # 占いする座標を20個ランダムに選ぶ
    S = set()
    while len(S)<20:
      i,j = randint(0,N-1),randint(0,N-1)
      S.add((i,j))
    
    # 占いをする
    print('q',len(S),end=' ')
    for i,j in S: print(i,j,end=' ')
    print(flush=True)
    r = int(input())
    
    for ind in range(n): 
      # i番目の盤面の時のv(s)
      vs = 0
      for i,j in S:
        vs += BOARD[ind][i][j]
      
      # 尤度をかける
      P[ind] *= q(len(S),vs,r)
      
    P = normalize(P)
    
    max_p = max(P)
    max_ind = P.index(max_p)
    
    # 確率が80%以上なら解答する
    if max_p>0.8:
      ANS = set()
      for i in range(N):
        for j in range(N):
          if BOARD[max_ind][i][j]>=1:
            ANS.add((i,j))
      
      print('a',len(ANS),end=' ')
      for i,j in ANS: print(i,j,end=' ')
      print(flush=True)
      r = int(input())
      if r==1: exit() # 正解なら終了
      else: P[max_ind] = 0.0 # 不正解なら事前確率を0にして終了
      
bayesian_inference(make_all_board())

このコードをseed=0で試してみると、cost=2.459675になりました。 占いをランダムにしているのにここまで改善するとは驚きです。

占いマス生成（相互情報量）

気が向いたら続きをやりたいと思います、、、