#859987 (Rust) No.5015 Escape from Labyrinth

• 問題: No.5015 Escape from Labyrinth
• ユーザー: ntk-ta01
• 提出日時: 2023-04-15 14:33:54
• 言語: Rust (1.83.0 + proconio)
• 結果: WA
• 実行時間: -
• コード長: 14,007 bytes
• コンパイル時間: 4,035 ms
• コンパイル使用メモリ: 163,288 KB
• 実行使用メモリ: 13,220 KB
• スコア: 6,470
• 最終ジャッジ日時: 2023-04-15 14:34:45
• 合計ジャッジ時間: 50,195 ms
• ジャッジサーバーID （参考情報）: judge14 / judge15

ファイルパターン	結果
other	AC * 6 WA * 94

ソースコード

// https://qiita.com/tanakh/items/0ba42c7ca36cd29d0ac8
macro_rules! input {
    (source = $s:expr, $($r:tt)*) => {
        let mut iter = $s.split_whitespace();
        input_inner!{iter, $($r)*}
    };
    ($($r:tt)*) => {
        let s = {
            use std::io::Read;
            let mut s = String::new();
            std::io::stdin().read_to_string(&mut s).unwrap();
            s
        };
        let mut iter = s.split_whitespace();
        input_inner!{iter, $($r)*}
    };
}

macro_rules! input_inner {
    ($iter:expr) => {};
    ($iter:expr, ) => {};

    ($iter:expr, $var:ident : $t:tt $($r:tt)*) => {
        let $var = read_value!($iter, $t);
        input_inner!{$iter $($r)*}
    };
}

macro_rules! read_value {
    ($iter:expr, ( $($t:tt),* )) => {
        ( $(read_value!($iter, $t)),* )
    };

    ($iter:expr, [ $t:tt ; $len:expr ]) => {
        (0..$len).map(|_| read_value!($iter, $t)).collect::<Vec<_>>()
    };

    ($iter:expr, chars) => {
        read_value!($iter, String).chars().collect::<Vec<char>>()
    };

    ($iter:expr, usize1) => {
        read_value!($iter, usize) - 1
    };

    ($iter:expr, $t:ty) => {
        $iter.next().unwrap().parse::<$t>().expect("Parse error")
    };
}

use std::{
    cell::RefCell,
    rc::{Rc, Weak},
};

const DIJ: [(usize, usize); 4] = [(0, !0), (!0, 0), (0, 1), (1, 0)];
const DIR: [char; 4] = ['L', 'U', 'R', 'D'];

fn main() {
    let (input, map) = parse_input();
    let out = beam_search(&input, &map);
    write_output(&out);
}

fn beam_search(input: &Input, map: &[Vec<Square>]) -> Output {
    const BEAM_WIDTH: usize = 1000;
    let mut tree = {
        let state = State::new(input, map.to_vec());
        let head = Rc::new(Node::new(None));
        BeamSearchTree { state, head }
    };
    let mut current_queue = vec![tree.head.clone()];
    let mut beam_queue = (0..input.h + 1).map(|_| vec![]).collect::<Vec<_>>();
    // let mut hash_set = FxHashSet::default();
    for turn in 0..input.h {
        tree.dfs(input, &mut beam_queue, turn, true);
        current_queue.clear();
        if turn + 1 == input.h {
            // 最終ターンなのでcandidatesを作らない
            break;
        }
        let mut candidates = vec![];
        std::mem::swap(&mut candidates, &mut beam_queue[turn + 1]);
        if candidates.len() > BEAM_WIDTH {
            candidates.select_nth_unstable_by_key(BEAM_WIDTH - 1, |c| std::cmp::Reverse(c.score));
            candidates.truncate(BEAM_WIDTH);
        }
        // hash_set.clear();
        for candidate in candidates {
            // プレイヤーの位置が同じ候補を重複除去
            // if !hash_set.insert(candidate.hash) {
            //     continue;
            // }
            let child = Node::new(Some(candidate.parent.clone()));
            let child_ptr = Rc::new(child);
            let children = &mut candidate.parent.children.borrow_mut();
            children.push((candidate.act, Rc::downgrade(&child_ptr)));
            current_queue.push(child_ptr);
        }
    }
    let last_candidates = beam_queue.pop().unwrap();
    let best_candidate = last_candidates
        .into_iter()
        .max_by_key(|state| state.score)
        .unwrap();

    // 復元
    let mut commands = vec![DIR[best_candidate.act.dir]];
    commands.reserve(input.h);
    let mut parent = best_candidate.parent;
    while let Some(p) = parent.parent.clone() {
        {
            let children = &mut p.children.borrow_mut();
            let (act, _) = children
                .iter()
                .find(|(_, child)| child.upgrade().is_some())
                .unwrap();
            commands.push(DIR[act.dir]);
        }
        parent = p;
    }
    commands.reverse();
    Output::new(commands)
}

#[derive(Clone)]
struct State {
    score: i32,
    map: Vec<Vec<Square>>,
    pos: (usize, usize),
    visited: Vec<Vec<bool>>,
    has_key: bool,
    turn: usize,
    // hash: u64,
}

impl State {
    fn new(input: &Input, map: Vec<Vec<Square>>) -> Self {
        let pos = find_player_position(&map);
        let mut visited = vec![vec![false; input.n]; input.n];
        visited[pos.0][pos.1] = true;
        State {
            score: 0,
            map,
            pos,
            visited,
            has_key: false,
            turn: 0,
            // hash,
        }
    }

    fn apply(&mut self, act: &Action) {
        self.turn += act.consume_turn;
        self.score += act.score_diff;
        // self.hash ^= act.hash_diff;

        let (player_r, player_c) = &mut self.pos;
        let next_r = *player_r + DIJ[act.dir].0;
        let next_c = *player_c + DIJ[act.dir].1;
        if act.is_moved {
            *player_r = next_r;
            *player_c = next_c;
        }
        if act.is_new_visited {
            self.visited[*player_r][*player_c] = true;
        }
        if act.is_got_key {
            self.has_key = true;
        }
    }

    fn revert(&mut self, act: &Action) {
        self.turn -= act.consume_turn;
        self.score -= act.score_diff;
        // self.hash ^= act.hash_diff;
        let rev_dir = (act.dir + 2) % 4;

        let (player_r, player_c) = &mut self.pos;
        let prev_r = *player_r + DIJ[rev_dir].0;
        let prev_c = *player_c + DIJ[rev_dir].1;
        if act.is_got_key {
            self.has_key = false;
        }
        if act.is_new_visited {
            self.visited[*player_r][*player_c] = false;
        }
        if act.is_moved {
            *player_r = prev_r;
            *player_c = prev_c;
        }
    }
}

struct Action {
    dir: usize,
    is_moved: bool,
    is_new_visited: bool,
    is_got_key: bool,
    score_diff: i32,
    consume_turn: usize,
    // hash_diff: u64,
}

impl Action {
    fn new(
        dir: usize,
        is_moved: bool,
        is_new_visited: bool,
        is_got_key: bool,
        score_diff: i32,
        consume_turn: usize,
        // hash_diff: u64,
    ) -> Self {
        Action {
            dir,
            is_moved,
            is_new_visited,
            is_got_key,
            score_diff,
            consume_turn,
            // hash_diff,
        }
    }
}

struct Candidate {
    score: i32,
    parent: Rc<Node>,
    act: Action,
    // hash: u64,
}

impl Candidate {
    // fn new(score: i32, parent: Rc<Node>, act: Action, hash: u64) -> Self {
    fn new(score: i32, parent: Rc<Node>, act: Action) -> Self {
        Candidate {
            score,
            parent,
            act,
            // hash,
        }
    }
}

struct Node {
    parent: Option<Rc<Node>>,
    children: RefCell<Vec<(Action, Weak<Node>)>>,
}

impl Node {
    fn new(parent: Option<Rc<Node>>) -> Self {
        let children = RefCell::new(vec![]);
        Self { parent, children }
    }
}

struct BeamSearchTree {
    state: State,
    head: Rc<Node>, // 木上の探索中のnodeを保持する
}

impl BeamSearchTree {
    fn dfs(
        &mut self,
        input: &Input,
        beam_queue: &mut Vec<Vec<Candidate>>,
        target_turn: usize,
        is_single_path: bool,
    ) {
        if self.state.turn == target_turn {
            for (dir, &(dr, dc)) in DIJ.iter().enumerate() {
                let mut score_diff = 0;
                let mut is_new_visited = false;
                let mut is_got_key = false;
                // let mut hash_diff = 0;

                let (player_r, player_c) = &self.state.pos;
                let next_r = *player_r + dr;
                let next_c = *player_c + dc;
                if input.n <= next_r || input.n <= next_c {
                    continue;
                }
                if self.state.map[next_r][next_c] == Square::Wall
                    || self.state.map[next_r][next_c] == Square::Finder
                {
                    continue;
                }
                let is_moved = true;
                if !self.state.visited[next_r][next_c] {
                    is_new_visited = true;
                    if self.state.map[next_r][next_c] == Square::Jewelry {
                        score_diff += 10;
                    }
                    if self.state.map[next_r][next_c] == Square::Key {
                        is_got_key = true;
                        score_diff += (input.h - self.state.turn) as i32 * 2;
                    }
                    if self.state.map[next_r][next_c] == Square::Door {
                        if self.state.has_key {
                            score_diff += (input.h - self.state.turn) as i32 * 2;
                        } else {
                            score_diff -= (input.h - self.state.turn) as i32 * 2;
                        }
                    }
                    // else if self.state.map[next_r][next_c] == Square::Finder {
                    // score_diff -= (input.t - self.state.turn) as i32 / 2;
                    // }
                }
                // hash_diff ^= input.player_hashes[*player_r][*player_c];
                // hash_diff ^= input.player_hashes[next_r][next_c];

                // next_turnが+1じゃないアクションが存在する問題のために
                // 下のif文を書いたり、beam_queueの長さをinput.t+1にする
                let consume_turn = || -> usize {
                    if self.state.map[next_r][next_c] == Square::Door {
                        return input.h - self.state.turn;
                    }
                    // 探知機にひっかかっていたらD減る
                    // TODO: 引っかかっているか調べるやつを書く
                    1
                }();
                let next_turn = self.state.turn + consume_turn;
                if next_turn < beam_queue.len() {
                    beam_queue[next_turn].push(Candidate::new(
                        self.state.score + score_diff,
                        self.head.clone(),
                        Action::new(
                            dir,
                            is_moved,
                            is_new_visited,
                            is_got_key,
                            score_diff,
                            consume_turn,
                        ),
                        // self.state.hash ^ hash_diff,
                    ));
                }
            }
        } else {
            let node = self.head.clone();

            let next_is_single_path = {
                let children = &mut node.children.borrow_mut();
                children.retain(|(_, child)| child.upgrade().is_some());
                if children.is_empty() {
                    false
                } else {
                    let next_is_single_path = is_single_path && (children.len() == 1);
                    for (act, child) in children.iter_mut() {
                        let next_turn = self.state.turn + 1;
                        if next_turn > target_turn {
                            continue;
                        }

                        self.head = child.upgrade().unwrap();
                        self.state.apply(act);

                        self.dfs(input, beam_queue, target_turn, next_is_single_path);

                        if !next_is_single_path {
                            self.state.revert(act);
                        }
                    }
                    next_is_single_path
                }
            };

            // nodeの子が1つでなければ、headを元のnodeに戻す
            if !next_is_single_path {
                self.head = node;
            }
        }
    }
}

#[derive(PartialEq, Eq, PartialOrd, Ord, Clone, Copy, Debug)]
enum Square {
    Empty,
    Wall,
    Player,
    Door,
    Key,
    Jewelry,
    Fire,
    Finder,
}

fn find_player_position(map: &[Vec<Square>]) -> (usize, usize) {
    for (r, row) in map.iter().enumerate() {
        for (c, square) in row.iter().enumerate() {
            if *square == Square::Player {
                return (r, c);
            }
        }
    }
    unreachable!("map上にplayerが見つかりませんでした")
}

fn parse_input() -> (Input, Vec<Vec<Square>>) {
    input! {
        n: usize,
        d: usize,
        h: usize,
        grid: [String; n],
        m: usize,
        yxd: [(usize, usize, usize); m],
    }
    (
        Input {
            n,
            d,
            h,
            m,
            finders: yxd
                .into_iter()
                .map(|(y, x, d)| Finder::new(y, x, d))
                .collect(),
        },
        grid.into_iter()
            .map(|row| {
                row.chars()
                    .map(|ch| match ch {
                        '.' => Square::Empty,
                        '#' => Square::Wall,
                        'S' => Square::Player,
                        'G' => Square::Door,
                        'K' => Square::Key,
                        'J' => Square::Jewelry,
                        'F' => Square::Fire,
                        'E' => Square::Finder,
                        _ => unreachable!("想定外の入力がありました"),
                    })
                    .collect()
            })
            .collect(),
    )
}

#[allow(dead_code)]
#[derive(Debug)]
struct Input {
    n: usize,             // グリッドの大きさ
    d: usize,             // 探知機によって削られる体力
    h: usize,             // プレイヤーの初期体力
    m: usize,             // 探知機の数
    finders: Vec<Finder>, // 探知機の位置と作動感覚
}

#[allow(dead_code)]
#[derive(Debug, Clone, Copy)]
struct Finder {
    y: usize, // y座標
    x: usize, // x座標
    d: usize, // 作動間隔
}

impl Finder {
    fn new(y: usize, x: usize, d: usize) -> Self {
        Finder { y, x, d }
    }
}

#[derive(Clone)]
struct Output {
    commands: Vec<char>,
}

impl Output {
    fn new(commands: Vec<char>) -> Self {
        Output { commands }
    }
}

fn write_output(out: &Output) {
    for cmd in out.commands.iter() {
        println!("M {}", cmd);
    }
}