ベイズ推定

[tipstar0125]

問題：https://atcoder.jp/contests/ahc030/tasks/ahc030_a 解答例：https://atcoder.jp/contests/ahc030/submissions/50496667

ベイズの定理より P(x|R)∝P(R|x)

x: 盤面の確率 R: 測定結果

測定結果Rのときの盤面xの確率を直接求めることは難しいが、 盤面xであったときの、測定結果Rの確率を求めることで、P(x|R)を推定できる。

参考1: https://qiita.com/aplysia/items/c3f2111110ac5043710a 参考1: https://bowwowforeach.hatenablog.com/entry/2023/08/24/205427?_gl=1*1rve9o5*_gcl_au*MTE5MTM3NjYzLjE2OTEzMjM4ODU

# https://qiita.com/aplysia/items/c3f2111110ac5043710a

import sys
from collections import deque
from copy import deepcopy
import random
from math import sqrt,erf

random.seed(0)
is_local=False # ローカルテスト時True
def eprint(s):print(s,file=sys.stderr)

# 入力
N,M,eps=list(map(float,input().split()))
N=int(N)
M=int(M)

class Coord:
    def __init__(self,y,x):
        self.y=y
        self.x=x

class Mino:
    def __init__(self,v):
        self.height=0
        self.width=0
        self.coords=[]
        for i in range(0,len(v),2):
            y=v[i]
            x=v[i+1]
            self.height=max(self.height,y+1)
            self.width=max(self.width,x+1)
            self.coords.append(Coord(y,x))

minos=[]
total=0
for _ in range(M):
    v=list(map(int,input().split()))
    minos.append(Mino(v[1:]))
    total+=v[0]

# ローカルテスト用入力(ローカル時コメントアウトを外す)
# ps=[]
# for _ in range(M):
#     y,x=list(map(int,input().split()))
#     ps.append(Coord(y,x))
# ans=[list(map(int,input().split())) for _ in range(N)]
# es=[float(input()) for _ in range(2*N*N)]

# 各ポリオミノを配置できる始点を求める
fixed_board=[[-1 for _ in range(N)] for _ in range(N)]  # 掘削して盤面を絞るときに使用（今は使用していない）(TODO)
candidate_mino_coords=[]
candidate_board_num=1
for mino in minos:
    cands=[]
    for y in range(N-mino.height+1):
        for x in range(N-mino.width+1):
            ok=True
            for coord in mino.coords:
                ny=y+coord.y
                nx=x+coord.x
                # 掘削済みで油田がなければ、置くことはできない(TODO)
                if fixed_board[ny][nx]==0:
                    ok=False
                    break
            if ok:cands.append(Coord(y,x))
    candidate_board_num*=len(cands)
    candidate_mino_coords.append(cands)

# 盤面のパターンが多い場合は一旦諦める
# 1マス掘削でパターン数を減らす(TODO)
if candidate_board_num>int(1e6):exit(0)

# BFSで盤面を生成
candidate_boards=[]
Q=deque()
Q.append((0,[[0 for _ in range(N)] for _ in range(N)]))
while Q:
    cnt,B=Q.popleft()
    if cnt==M:
        ok=True
        for y in range(N):
            for x in range(N):
                # 生成した盤面の油田数が既に掘削した油田数と一致していなければ、候補から外す(TODO)
                if fixed_board[y][x]!=-1 and fixed_board[y][x]!=B[y][x]:
                    ok=False
                    break
            if not ok:break
        if ok:candidate_boards.append(B)
        continue

    for st in candidate_mino_coords[cnt]:
        nB=deepcopy(B)
        ok=True
        for mino in minos[cnt].coords:
            ny=st.y+mino.y
            nx=st.x+mino.x
            nB[ny][nx]+=1
            # 生成途中の盤面の油田数が既に掘削した油田数より多くなたら候補から外す(TODO)
            if fixed_board[y][x]!=-1 and fixed_board[y][x]<B[y][x]:
                ok=False
                break
            if not ok:break
        if ok:Q.append((cnt+1,nB))

# クエリパート
turn=0
cost=0.0

def query(coords):
    global turn
    global cost
    k=len(coords)
    turn+=1
    cost+=1/sqrt(k)

    def make_query_return():
        if not is_local:return int(input())
        if k==1:
            return ans[coords[0].y][coords[0].x]
        else:
            vs=0.0
            for coord in coords:
                vs+=ans[coord.y][coord.x]
            mean=(k-vs)*eps+vs*(1.0-eps)
            std=sqrt(k*eps*(1.0-eps))
            ret=mean+std*es[turn]
            ret=int(round(ret))
            return max(0,ret)

    q=["q",k]
    for coord in coords:
        q.append(coord.y)
        q.append(coord.x)
    print(*q,flush=True)

    ret=make_query_return()
    return ret

# アンサーパート
def answer(coords):
    k=len(coords)

    def make_answer_return():
        if not is_local:return int(input())
        ok=True
        cnt=0
        for coord in coords:
            c=ans[coord.y][coord.x]
            if c==0:
                ok=False
                break
            else:cnt+=c
        ok&=cnt==total
        if ok:return 1
        else:return 0

    a=["a",k]
    for coord in coords:
        a.append(coord.y)
        a.append(coord.x)
    print(*a,flush=True)

    ret=make_answer_return()
    return ret

# 尤度計算
# https://bowwowforeach.hatenablog.com/entry/2023/08/24/205427?_gl=1*1rve9o5*_gcl_au*MTE5MTM3NjYzLjE2OTEzMjM4ODU
def likelihood(k,cnt,ret):
    mean=(k-cnt)*eps+cnt*(1.0-eps)
    std=sqrt(k*eps*(1.0-eps))
    diff=ret-mean

    def prob_in_range(l,r):
        def cdf(x):return 0.5*(1.0+erf(x)/(std*sqrt(2)))
        return cdf(r)-cdf(l)

    if ret==0:return prob_in_range(-1e10,diff+0.5)
    else:return prob_in_range(diff-0.5,diff+0.5)

# 規格化
def normalize(prob):
    s=sum(prob)
    for i in range(len(prob)):
        prob[i]/=s
    return prob

# ベイズ推定パート
board_num=len(candidate_boards)
prob=[1/board_num for _ in range(board_num)] # 同確率で初期化

while turn<2*N*N:
    k=30 # 工夫できそう。情報量最大化とか(TODO)
    coords=[]
    st=set()
    while len(st)<k:
        y=random.randrange(N)
        x=random.randrange(N)
        if (y,x) in st:continue
        st.add((y,x))
        coords.append(Coord(y,x))
    ret=query(coords)

    for i in range(board_num):
        cnt=0
        for coord in coords:
            cnt+=candidate_boards[i][coord.y][coord.x]
        # 確率がめちゃくちゃ小さくなって、まともに計算できない場合がありそう
        # 規格化で0割りもありえそうなので、対数を取った方がよさそう(TODO)
        prob[i]*=likelihood(k,cnt,ret)
    prob=normalize(prob)

    mx=max(prob)
    idx=prob.index(mx)

    if mx>0.8: # 工夫できそう。第二候補との差とか(TODO)
        a=[]
        for y in range(N):
            for x in range(N):
                if candidate_boards[idx][y][x]>0:a.append(Coord(y,x))
        ret=answer(a)
        if ret==1:
            eprint(f"Turn: {turn}")
            eprint(f"Cost: {cost}")
            exit(0)
        else:
            cost+=1.0
            prob[idx]=0.0

[tipstar0125]

相互情報量

参考：https://kiri8128.hatenablog.com/entry/2024/02/19/211740

事前分布P(x)、事後分布Q(x)としたときの情報量は下式で表される。

今までの計測（占い）によって得られた分布より、任意のkを選んだときの油田の数を求める。 これは、各候補となる盤面に対して(盤面の確率)x(kマスの油田数)を求めて、平均した値と期待できる。 この値がクエリを投げた結果の戻り値だと仮定して、ベイズで仮の事後分布を求めて、情報量を求める式より、値を得る。 これをコストで割った値が最大になるものを選ぶことで、コストを抑えることができる。

ただし、kの数やパターンを考えると計算が間に合わないので、パターンを絞る工夫が必要。

• kは固定
• 隣同士の方が良さそう
• 一度測定した箇所から得られる情報量は多くない

【結論】 なんか違いそう。。。 もう少し式レベルで考えないといけない。

[tipstar0125]

焼き鈍し

参考：https://eijirou-kyopro.hatenablog.com/entry/2024/02/22/152604

[tipstar0125]

その他参考 terryさん：https://www.terry-u16.net/entry/ahc030 あぷりしあさん（相互情報量）：https://qiita.com/aplysia/items/29a4fb4573fc1b8dec79?utm_campaign=post_article&utm_medium=twitter&utm_source=twitter_share AHCラジオ：https://www.youtube.com/watch?v=YvCYsiu-TQs&t=6s Write解：https://atcoder.jp/contests/ahc030/submissions/50443474

[tipstar0125]

M=2 eps=0.2のときに10秒かかる。。。 ※最初は固定占いで相互情報量の多いものを選び、3bitより小さくなったら、山登り占い。

#![allow(non_snake_case)]
#![allow(unused_imports)]
#![allow(unused_macros)]
#![allow(clippy::comparison_chain)]
#![allow(clippy::nonminimal_bool)]
#![allow(clippy::neg_multiply)]
#![allow(clippy::type_complexity)]
#![allow(clippy::needless_range_loop)]
#![allow(dead_code)]

use std::{
    cmp::Reverse,
    collections::{BTreeMap, BTreeSet, BinaryHeap, HashSet, VecDeque},
};

use im_rc::HashMap;
use itertools::Itertools;
use proconio::{
    input,
    marker::{Chars, Usize1},
};
use rand::prelude::*;
use superslice::Ext;

fn main() {
    let start = std::time::Instant::now();

    let mut stdin =
        proconio::source::line::LineSource::new(std::io::BufReader::new(std::io::stdin()));
    macro_rules! input(($($tt:tt)*) => (proconio::input!(from &mut stdin, $($tt)*)));

    input! {
        N: usize,
        M: usize,
        eps: f64,
    }

    let mut minos = vec![];
    for _ in 0..M {
        input! {
            d: usize,
            coords: [(usize, usize); d]
        }
        let mut coord_diff = vec![];
        let mut height = 0;
        let mut width = 0;
        for coord in coords.iter() {
            let row = coord.0;
            let col = coord.1;
            height = height.max(row + 1);
            width = width.max(col + 1);
            coord_diff.push(CoordDiff::new(row as isize, col as isize));
        }
        let mino = Mino {
            coord_diff,
            height,
            width,
        };
        minos.push(mino);
    }

    // Local
    input! {
        _ps: [(usize, usize); M],
        ans: [[i16; N]; N],
        es: [f64; 2*N*N]
    }
    let ans = ans.into_iter().flatten().collect_vec();
    let ans = DynamicMap2d::new(ans, N);

    let candidate_mino_coords = make_candidate_mino_coords(N, &minos);
    let boards = make_boards(N, &minos, &candidate_mino_coords);

    let mut mino_num = 0;
    for mino in minos.iter() {
        mino_num += mino.coord_diff.len();
    }

    let mut turn = 0;
    let mut cost = 0.0_f64;
    let mut rng = rand_chacha::ChaCha20Rng::seed_from_u64(1);

    let board_num = boards.len();
    let mut prob = vec![1.0 / (board_num as f64); board_num];

    let mut log2 = vec![];
    for i in 0..=1e4 as usize + 1 {
        let x = i as f64 / 1e4;
        log2.push(x.log2());
    }

    let climbing_iteration_num = 1000;
    let mut fixed_fortune_coords = vec![];

    for sep_num in 2..3 {
        for i in 0..N {
            if i + N / sep_num >= N {
                continue;
            }
            for j in 0..N {
                if j + N / sep_num >= N {
                    continue;
                }
                let mut coords = vec![];
                for k in 0..N / sep_num {
                    for l in 0..N / sep_num {
                        coords.push(Coord::new(i + k, j + l));
                    }
                }
                fixed_fortune_coords.push((coords.len(), coords));
            }
        }
    }

    eprintln!("fixed_fortune_num: {}", fixed_fortune_coords.len());
    let info_amount_threshold = 3.0;
    while !fixed_fortune_coords.is_empty() {
        let mut best_score = 0.0;
        let mut best_k = 0;
        let mut best_coords = vec![];
        let mut remove_index = vec![];

        for i in 0..fixed_fortune_coords.len() {
            let mut cnt = vec![0_i16; board_num];
            let (k, coords) = fixed_fortune_coords[i].clone();
            for (c, b) in cnt.iter_mut().zip(boards.iter()) {
                for coord in coords.iter() {
                    *c += b[*coord];
                }
            }
            let score = calc_mutual_information(k, eps, &prob, &cnt, mino_num, &log2);
            if score < info_amount_threshold {
                remove_index.push(i);
            }
            if score > best_score {
                best_score = score;
                best_k = k;
                best_coords = coords.clone();
            }
        }
        let k = best_k;
        let coords = best_coords;
        if best_score < info_amount_threshold {
            break;
        }
        remove_index.sort();
        remove_index.reverse();
        for idx in remove_index.iter() {
            fixed_fortune_coords.remove(*idx);
        }

        let ret = query(&coords, &ans, eps, &es, &mut turn, &mut cost);
        make_query(&coords);
        // input! {ret:u8};
        for i in 0..board_num {
            let mut cnt = 0;
            for coord in coords.iter() {
                cnt += boards[i][*coord];
            }
            prob[i] *= likelihood(k, eps, cnt, ret);
        }
        normalize(&mut prob);
        let mut p_cnt = 0;
        for p in prob.iter() {
            if *p > 1e-4 {
                p_cnt += 1;
            }
        }
        if p_cnt * climbing_iteration_num < 1e5 as usize {
            break;
        }
    }
    eprintln!("fixed_fortune_num: {}", fixed_fortune_coords.len());

    loop {
        let coords = decide_fortune(
            N,
            eps,
            &mut rng,
            climbing_iteration_num,
            &prob,
            &boards,
            mino_num,
            &log2,
        );
        let k = coords.len();
        // eprintln!("{}", k);
        let ret = query(&coords, &ans, eps, &es, &mut turn, &mut cost);
        make_query(&coords);
        // input! {ret:u8};
        for i in 0..board_num {
            let mut cnt = 0;
            for coord in coords.iter() {
                cnt += boards[i][*coord];
            }
            prob[i] *= likelihood(k, eps, cnt, ret);
        }
        normalize(&mut prob);
        let mut mx = 0.0;
        let mut idx = 0;
        for (i, &p) in prob.iter().enumerate() {
            if mx < p {
                mx = p;
                idx = i;
            }
        }
        // let mut pp = vec![];
        // for (i, p) in prob.iter().enumerate() {
        //     if *p > 0.01 {
        //         pp.push((i, *p));
        //     }
        // }
        // eprintln!("{:?}", pp);
        if mx > 0.8 {
            make_answer(&boards[idx]);
            let ret = eval(&boards[idx], &ans);
            eprintln!("{idx}");
            // input! {ret:u8};
            if ret == 1 {
                break;
            }
            cost += 1.0;
            prob[idx] = 0.0;
        }
    }

    let score = (1e6 * cost.max(1.0 / N as f64)).round() as usize;
    eprintln!("Turn: {}", turn as f64 / (2.0 * N as f64 * N as f64));
    eprintln!("Cost: {}", cost);
    eprintln!("Score: {score}");

    #[allow(unused_mut, unused_assignments)]
    let mut elapsed_time = start.elapsed().as_micros() as f64 * 1e-6;
    #[cfg(feature = "local")]
    {
        eprintln!("Local Mode");
        elapsed_time *= 0.55;
    }
    eprintln!("Elapsed: {}", (elapsed_time * 1000.0) as usize);
}

fn make_candidate_mino_coords(N: usize, minos: &[Mino]) -> Vec<Vec<Coord>> {
    let mut ret = vec![];
    for mino in minos.iter() {
        let mut cands = vec![];
        for i in 0..N - mino.height + 1 {
            for j in 0..N - mino.width + 1 {
                let pos = Coord::new(i, j);
                cands.push(pos);
            }
        }
        ret.push(cands);
    }
    ret
}

fn make_boards(
    N: usize,
    minos: &[Mino],
    candidate_mino_coords: &[Vec<Coord>],
) -> Vec<DynamicMap2d<i16>> {
    let M = minos.len();
    let mut ret = vec![];
    let mut Q = VecDeque::new();
    Q.push_back((0, DynamicMap2d::new(vec![0_i16; N * N], N)));
    while let Some((cnt, B)) = Q.pop_front() {
        if cnt == M {
            ret.push(B.clone());
            continue;
        }
        for &pos in candidate_mino_coords[cnt].iter() {
            let mut nB = B.clone();
            for &coord_diff in minos[cnt].coord_diff.iter() {
                let nxt = pos + coord_diff;
                nB[nxt] += 1;
            }
            Q.push_back((cnt + 1, nB));
        }
    }
    ret
}

#[allow(clippy::too_many_arguments)]
fn decide_fortune(
    N: usize,
    eps: f64,
    rng: &mut rand_chacha::ChaCha20Rng,
    iteration_num: usize,
    prob: &[f64],
    boards: &[DynamicMap2d<i16>],
    mino_num: usize,
    log2: &[f64],
) -> Vec<Coord> {
    let mut fortune_map = DynamicMap2d::new(vec![true; N * N], N);
    let mut best_score = 0.0;
    let mut k = (N * N) as i16;

    let L = boards.len();
    let mut cnt = vec![0_i16; L];
    for (c, b) in cnt.iter_mut().zip(boards) {
        for i in 0..N {
            for j in 0..N {
                let coord = Coord::new(i, j);
                if fortune_map[coord] {
                    *c += b[coord];
                }
            }
        }
    }

    for _ in 0..iteration_num {
        let i = rng.gen_range(0, N);
        let j = rng.gen_range(0, N);
        let coord = Coord::new(i, j);

        let delta = if fortune_map[coord] { -1 } else { 1 };
        fortune_map[coord] = !fortune_map[coord];
        for (i, board) in boards.iter().enumerate() {
            cnt[i] += board[coord] * delta;
        }
        k += delta;

        let score = calc_mutual_information(k as usize, eps, prob, &cnt, mino_num, log2);
        if score > best_score {
            best_score = score;
        } else {
            let delta = if fortune_map[coord] { -1 } else { 1 };
            fortune_map[coord] = !fortune_map[coord];
            for (i, board) in boards.iter().enumerate() {
                cnt[i] += board[coord] * delta;
            }
            k += delta;
        }
    }
    let mut fortune_coord = vec![];
    for i in 0..N {
        for j in 0..N {
            let coord = Coord::new(i, j);
            if fortune_map[coord] {
                fortune_coord.push(coord);
            }
        }
    }
    fortune_coord
}

fn calc_mutual_information(
    k: usize,
    eps: f64,
    prob: &[f64],
    cnt: &[i16],
    mino_num: usize,
    log2: &[f64],
) -> f64 {
    let mean_max = (k as f64 - mino_num as f64) * eps + (mino_num as f64) * (1.0 - eps);
    let std_dev = ((k as f64) * eps * (1.0 - eps)).sqrt();
    let query_result_num = (mean_max + 3.0 * std_dev) as usize;
    let mut query_result_p_vec = vec![0.0; query_result_num + 1];
    let prune_threshold = 1e-4; // <0.01%

    for (&p, &c) in prob.iter().zip(cnt) {
        if p < prune_threshold {
            continue;
        }
        let mean = (k as f64 - c as f64) * eps + (c as f64) * (1.0 - eps);
        let lower = (mean - 3.0 * std_dev).max(0.0) as usize;
        let upper = ((mean + 3.0 * std_dev) as usize).min(query_result_num);
        for query_result in lower..=upper {
            query_result_p_vec[query_result] += p * likelihood(k, eps, c, query_result);
        }
    }
    let mut ret = 0.0;

    for (query_result, query_result_p) in query_result_p_vec.iter().enumerate() {
        if *query_result_p < prune_threshold {
            continue;
        }
        for (&p, &c) in prob.iter().zip(cnt) {
            if p < prune_threshold {
                continue;
            }
            let likelihood = likelihood(k, eps, c, query_result);
            if likelihood < prune_threshold {
                continue;
            }
            let likelihood_usize = (likelihood * 1e4) as usize;
            let query_result_p_usize = (query_result_p * 1e4) as usize;
            ret += likelihood * p * (log2[likelihood_usize] - log2[query_result_p_usize]);
        }
    }

    ret * (k as f64).sqrt()
}

#[allow(clippy::approx_constant)]
fn normal_cdf(x: f64, mean: f64, std_dev: f64) -> f64 {
    0.5 * (1.0 + libm::erf((x - mean) / (std_dev * 1.41421356237)))
}

fn probability_in_range(mean: f64, std_dev: f64, a: f64, b: f64) -> f64 {
    if mean < a {
        return probability_in_range(mean, std_dev, 2.0 * mean - b, 2.0 * mean - a);
    }
    let p_a = normal_cdf(a, mean, std_dev);
    let p_b = normal_cdf(b, mean, std_dev);
    p_b - p_a
}

fn likelihood(k: usize, eps: f64, cnt: i16, res: usize) -> f64 {
    let mean = (k as f64 - cnt as f64) * eps + (cnt as f64) * (1.0 - eps);
    let std_dev = ((k as f64) * eps * (1.0 - eps)).sqrt();
    if res == 0 {
        probability_in_range(mean, std_dev, -1e10, res as f64 + 0.5)
    } else {
        probability_in_range(mean, std_dev, res as f64 - 0.5, res as f64 + 0.5)
    }
}

fn normalize(prob: &mut [f64]) {
    let s = prob.iter().sum::<f64>();
    assert!(s >= 0.0);
    for p in prob.iter_mut() {
        *p /= s;
    }
}

#[derive(Debug)]
struct Mino {
    coord_diff: Vec<CoordDiff>,
    height: usize,
    width: usize,
}

fn query(
    coords: &[Coord],
    ans: &DynamicMap2d<i16>,
    eps: f64,
    es: &[f64],
    turn: &mut usize,
    cost: &mut f64,
) -> usize {
    let k = coords.len() as f64;
    *turn += 1;
    *cost += 1.0 / k.sqrt();
    if k == 1.0 {
        return ans[coords[0]] as usize;
    }
    let mut vs = 0.0;
    for coord in coords.iter() {
        vs += ans[*coord] as f64;
    }
    let mean = (k - vs) * eps + vs * (1.0 - eps);
    let std = (k * eps * (1.0 - eps)).sqrt();
    let ret = mean + std * es[*turn];
    let ret = ret.round() as usize;
    ret.max(0)
}

fn make_query(coords: &[Coord]) {
    let mut v = vec![];
    v.push(coords.len());
    for coord in coords.iter() {
        v.push(coord.row);
        v.push(coord.col);
    }
    let mut query = "q ".to_string();
    query += v.iter().join(" ").as_str();
    println!("{query}");
}

fn make_answer(a: &DynamicMap2d<i16>) {
    let N = a.size;
    let mut coords = vec![];
    for i in 0..N {
        for j in 0..N {
            let coord = Coord::new(i, j);
            if a[coord] > 0 {
                coords.push(coord);
            }
        }
    }
    let mut v = vec![];
    v.push(coords.len());
    for coord in coords.iter() {
        v.push(coord.row);
        v.push(coord.col);
    }
    let mut ans = "a ".to_string();
    ans += v.iter().join(" ").as_str();
    println!("{ans}");
}

fn eval(a: &DynamicMap2d<i16>, ans: &DynamicMap2d<i16>) -> u8 {
    let N = a.size;
    for i in 0..N {
        for j in 0..N {
            let pos = Coord::new(i, j);
            if ans[pos] != a[pos] {
                return 0;
            }
        }
    }
    1
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
pub struct Coord {
    row: usize,
    col: usize,
}

impl Coord {
    pub fn new(row: usize, col: usize) -> Self {
        Self { row, col }
    }
    pub fn in_map(&self, height: usize, width: usize) -> bool {
        self.row < height && self.col < width
    }
    pub fn to_index(&self, width: usize) -> CoordIndex {
        CoordIndex(self.row * width + self.col)
    }
}

impl std::ops::Add<CoordDiff> for Coord {
    type Output = Coord;
    fn add(self, rhs: CoordDiff) -> Self::Output {
        Coord::new(
            self.row.wrapping_add_signed(rhs.dr),
            self.col.wrapping_add_signed(rhs.dc),
        )
    }
}

#[derive(Debug, Clone, Copy)]
pub struct CoordDiff {
    dr: isize,
    dc: isize,
}

impl CoordDiff {
    pub const fn new(dr: isize, dc: isize) -> Self {
        Self { dr, dc }
    }
}

pub const ADJ: [CoordDiff; 4] = [
    CoordDiff::new(1, 0),
    CoordDiff::new(!0, 0),
    CoordDiff::new(0, 1),
    CoordDiff::new(0, !0),
];

pub struct CoordIndex(pub usize);

impl CoordIndex {
    pub fn new(index: usize) -> Self {
        Self(index)
    }
    pub fn to_coord(&self, width: usize) -> Coord {
        Coord {
            row: self.0 / width,
            col: self.0 % width,
        }
    }
}

#[derive(Debug, Clone)]
pub struct DynamicMap2d<T> {
    pub size: usize,
    map: Vec<T>,
}

impl<T> DynamicMap2d<T> {
    pub fn new(map: Vec<T>, size: usize) -> Self {
        assert_eq!(size * size, map.len());
        Self { size, map }
    }
}

impl<T> std::ops::Index<Coord> for DynamicMap2d<T> {
    type Output = T;

    #[inline]
    fn index(&self, coordinate: Coord) -> &Self::Output {
        &self[coordinate.to_index(self.size)]
    }
}

impl<T> std::ops::IndexMut<Coord> for DynamicMap2d<T> {
    #[inline]
    fn index_mut(&mut self, coordinate: Coord) -> &mut Self::Output {
        let size = self.size;
        &mut self[coordinate.to_index(size)]
    }
}

impl<T> std::ops::Index<CoordIndex> for DynamicMap2d<T> {
    type Output = T;

    fn index(&self, index: CoordIndex) -> &Self::Output {
        unsafe { self.map.get_unchecked(index.0) }
    }
}

impl<T> std::ops::IndexMut<CoordIndex> for DynamicMap2d<T> {
    #[inline]
    fn index_mut(&mut self, index: CoordIndex) -> &mut Self::Output {
        unsafe { self.map.get_unchecked_mut(index.0) }
    }
}

[tipstar0125]

Write解：https://atcoder.jp/contests/ahc030/submissions/50443474

高速化の工夫

• Sim(new)
  • 全てのk(1～N^2)、占い範囲で真値としてありえる油田数の総和(cnt(B))、想定されるクエリ結果(res)、に対して、尤度P(res|B)を事前計算
  • 対数も同時に計算しておく
• Query(eval)
  • ↑の事前計算を用いる
  • log(x/y)の計算はlog(x)-log(y)として、logの計算は事前計算したものを用いて、ループ内では引き算のみとする
• クエリの作り方(main)
  • 各マスの相互情報量を求めて、大きい順にクエリの使用有無をフリップ
  • 配置候補で埋蔵量が同じであるマス（相互情報量がゼロ）はクエリから除外
  • （高速化ではないが）情報量0のマスでも、得られたクエリに追加することで、コスト下げられる
  • ループは3回！？

配置候補の尤度計算

• 配置候補はたくさんあるので、M=2以外は固定にできない
• 配置候補が変化するので、前計算を更新していくことができないので、クエリ毎にあり得る尤度計算を事前にしておき、尤度が必要になったときに、都度クエリの計算を最初から行う
• 対数尤度の時は足し算
• 規格化のときは、対数尤度から戻して計算

[tipstar0125]

ローカルテスタ、コマンド

cargo run -r --manifest-path tools/Cargo.toml --bin tester cargo run -r --bin c < in > out