AtCoder Beginner Contest 361 D

https://atcoder.jp/contests/abc361/tasks/abc361_d

各マスはWかBか空なので2ビットで表されます。あとはしらみつぶしで。

// Go Stone Puzzle
#![allow(non_snake_case)]

use std::collections::{VecDeque, HashSet};


//////////////////// library ////////////////////

fn read<T: std::str::FromStr>() -> T {
    let mut line = String::new();
    std::io::stdin().read_line(&mut line).ok();
    line.trim().parse().ok().unwrap()
}

fn read_vec<T: std::str::FromStr>() -> Vec<T> {
    read::<String>().split_whitespace()
            .map(|e| e.parse().ok().unwrap()).collect()
}


//////////////////// process ////////////////////

fn read_input() -> (usize, String, String) {
    let N: usize = read();
    let S: String = read();
    let T: String = read();
    (N, S, T)
}

type State = u32;

fn str_to_quat(S: String) -> State {
    let mut n: u32 = 15;    // 右端の2マス
    for c in S.chars().rev() {
        n = n * 4 + (if c == 'W' { 0 } else { 1 })
    }
    n
}

fn spaces(s: State, N: usize) -> (usize, usize) {
    let mut i1 = 100;
    let mut i2 = 0;
    for i in 0..N+2 {
        if ((s >> ((i as usize)*2)) & 3) == 3 {
            if i1 == 100 {
                i1 = i
            }
            else {
                i2 = i
            }
        }
    }
    (i1, i2)
}

fn get_cell(i: usize, s: State) -> u32 {
    (s >> ((i as u32) * 2)) & 3
}

fn set_cell(i: usize, s1: u32, s: State) -> State {
    let j = (i as u32) * 2;
    s ^ (((s >> j) & 3) << j) | (s1 << j)
}

fn nexts(s0: State, N: usize) -> Vec<State> {
    let mut states: Vec<State> = vec![];
    let (i1, i2) = spaces(s0, N);
    for i in 0..N+1 {
        if i == i1 || i == i2 || i+1 == i1 || i+1 == i2 {
            continue
        }
        let mut s = s0;
        let s1 = get_cell(i, s0);
        let s2 = get_cell(i+1, s0);
        s = set_cell(i1, s1, s);
        s = set_cell(i2, s2, s);
        s = set_cell(i, 3, s);
        s = set_cell(i+1, 3, s);
        states.push(s)
    }
    states
}

fn F(N: usize, S: String, T: String) -> i32 {
    if S == T {
        return 0
    }
    
    let start = str_to_quat(S);
    let end = str_to_quat(T);
    let mut visited: HashSet<State> = HashSet::new();
    let mut queue: VecDeque<(i32, State)> = VecDeque::new();
    queue.push_back((0, start));
    visited.insert(start);
    while let Some((n, s)) = queue.pop_front() {
        let ss = nexts(s, N);
        for s1 in ss.into_iter() {
            if s1 == end {
                return n+1
            }
            else if !visited.contains(&s1) {
                queue.push_back((n+1, s1));
                visited.insert(s1);
            }
        }
    }
    -1
}

fn main() {
    let (N, S, T) = read_input();
    println!("{}", F(N, S, T))
}