LeetCode Rust solution for the Sliding Puzzle problem #773

Question

Here is the solution to the LeetCode Sliding Puzzle Problem #773. The problem is to find the shortest solution for a 3 x 2 sliding puzzle.

I would appreciate any ideas to improve this code. But it's important not to break the performance.

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

impl Solution {
    pub fn sliding_puzzle(board: Vec<Vec<i32>>) -> i32 {
        // Define the final board state and its ID.
        let final_tiles = vec![1, 2, 3, 4, 5, 0];
        let final_id = Solution::board_id(&final_tiles);
        let final_zpos = 5; // Empty tile in the final board is always at index 5.

        // Convert the input board to a flat representation and its ID.
        let mut start_tiles = Vec::new();
        for row in &board {
            start_tiles.extend(row);
        }
        let start_id = Solution::board_id(&start_tiles);

        // Directions for movement: left, right, down, up.
        let directions = vec![(0, -1), (0, 1), (1, 0), (-1, 0)];

        // BFS initialization.
        let mut queue = VecDeque::new();
        queue.push_back((final_id, final_zpos));
        let mut dist = HashMap::new();
        dist.insert(final_id, 0);

        // Perform BFS.
        while let Some((current_id, zpos)) = queue.pop_front() {
            let current_dist = dist[&current_id];
            let mut tiles = Solution::from_id(current_id);

            for &(di, dj) in &directions {
                if let Some((next_zpos, next_id)) = Solution::do_step(&mut tiles, zpos, di, dj) {
                    if !dist.contains_key(&next_id) {
                        dist.insert(next_id, current_dist + 1);
                        queue.push_back((next_id, next_zpos));
                    }
                    // Undo the move to restore the tiles.
                    Solution::undo_step(&mut tiles, next_zpos, zpos);
                }
            }
        }

        // Return the result.
        dist.get(&start_id).cloned().unwrap_or(-1)
    }

    // Convert board state to a unique ID.
    fn board_id(tiles: &[i32]) -> i32 {
        tiles.iter().fold(0, |id, &t| id * 6 + t)
    }

    // Convert a unique ID back to board tiles.
    fn from_id(mut id: i32) -> Vec<i32> {
        let mut tiles = vec![0; 6];
        for i in (0..6).rev() {
            tiles[i] = id % 6;
            id /= 6;
        }
        tiles
    }

    // Perform a move if valid and return the new zero position and board ID.
    fn do_step(tiles: &mut Vec<i32>, zpos: usize, di: i32, dj: i32) -> Option<(usize, i32)> {
        let r = zpos as i32 / 3;
        let c = zpos as i32 % 3;
        let nr = r + di;
        let nc = c + dj;

        if nr >= 0 && nr < 2 && nc >= 0 && nc < 3 {
            let new_zpos = (nr * 3 + nc) as usize;
            tiles.swap(zpos, new_zpos);
            Some((new_zpos, Solution::board_id(tiles)))
        } else {
            None
        }
    }

    // Undo a move by swapping back the tiles.
    fn undo_step(tiles: &mut Vec<i32>, zpos: usize, prev_zpos: usize) {
        tiles.swap(zpos, prev_zpos);
    }
}

Answer 1

First of all, this code is good: it achieves its intended purpose, is readable and understandable and, within the leetcode restrictions, is reasonably architectured. Bravo: this is way better than a lot of stuff I have to work with!

Let's start with the most trivial:

Clippy concerns

Just running clippy on your snippet, I got a few warnings:

• usage of contains_key followed by insert on a HashMap
• writing &mut Vec instead of &mut [_] involves a new object where a slice will do
• manual Range::contains implementation

The first one is performance-related: try to avoid lookup-then-insert flows if possible. The second one is just a "best practice": there's no need to restrict your inputs to Vec where a slice can work too. The last warning is purely stylistic AFAIC: instead of 0 <= x && x < X you can check if (0..X).contains(x) which reads a little better.

Testing

I understand that Leetcode provides you with tests, but they do not exist outside of Leetcode. Since their editor is in browser and, well, borderline unusable, it would be nice to have a copy&paste test suite.

Style

The whole snippet is built around primitive types - mostly i32. You often refer to a tuple (board_id, zpos) - it would be nice as a separate struct, though it's usually fine for competitive programming.

Performance

You mentioned performance in your question. I don't think it's really relevant here (the problem space is too small, your code runs in 40ms on my PC), but let's look at a few low-hanging optimizations:

• early return. This is the cheapest one: if you encountered a solution during BFS, you have no chance to find any better. Just stop there.
• eliminate board-number and number-board conversions. You have a nice idea of representing the board as a single integer, and what you call an "id" is in fact just a compact representation. Use that! Swapping two cells is as easy as doing old_board + moved_value * 6**(5 - old_zero_position) - moved_value * 6**(5 - new_zero_position).
• You chose 6 as a "step" between cells for ID generation. If you use 8 instead, the number still fits in u32, while the number crunching becomes more computer friendly because 8 is a power of two and can be replaced by bit shifts).
• To save a lookup, replace HashMap with HashSet and push the distance to the queue together with a new board.

Suggestions

I tried to rewrite your solution to improve the general architecture and performance (benchmarks will follow).

My rust isn't great, maybe I'll post this as a review question too - there should be a lot to improve here.

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

pub struct Solution;
impl Solution {
    pub fn sliding_puzzle(board: Vec<Vec<i32>>) -> i32 {
        let final_board = Board::from_vec(&[1, 2, 3, 4, 5, 0]);
        let start_board = Board::from_leetcode(&board);

        if start_board == final_board {
            return 0;
        }

        // BFS initialization.
        let mut queue = VecDeque::with_capacity(511);
        queue.push_back((final_board, 0));
        let mut visited = HashSet::with_capacity(511);

        // Perform BFS.
        while let Some((current_board, dist)) = queue.pop_front() {
            for dir in &Direction::OPTIONS {
                let Some(next_board) = current_board.move_zero(dir) else {
                    continue;
                };
                if visited.insert(next_board.clone()) {
                    if next_board == start_board {
                        return dist + 1;
                    }
                    queue.push_back((next_board, dist + 1));
                }
            }
        }

        // Not found
        -1
    }
}

#[derive(Clone, Debug)]
enum Direction {
    Left,
    Right,
    Up,
    Down,
}

impl Direction {
    pub const OPTIONS: [Self; 4] = [Self::Left, Self::Right, Self::Up, Self::Down];

    pub const fn as_offset(&self) -> (i32, i32) {
        match self {
            Self::Left => (0, -1),
            Self::Right => (0, 1),
            Self::Up => (-1, 0),
            Self::Down => (1, 0),
        }
    }
}

#[derive(Clone, Debug, Eq)]
struct Board {
    /// Board state encoded as Self::BITS_PER_TILE-bit groups.
    state: u32,
    /// Zero position. Internal field for performance: can be derived from state.
    zpos: u32,
}

impl Board {
    const WIDTH: u32 = 3;
    const HEIGHT: u32 = 2;
    const SIZE: u32 = Self::WIDTH * Self::HEIGHT;
    const BITS_PER_TILE: u32 = 4;
    const TILE_MASK: u32 = (1 << Self::BITS_PER_TILE) - 1;

    pub fn from_leetcode(board: &[Vec<i32>]) -> Self {
        let mut tiles: Vec<i32> = Vec::with_capacity(6);
        for row in board {
            tiles.extend(row);
        }
        Self::from_vec(&tiles)
    }

    pub fn from_vec(tiles: &[i32]) -> Self {
        debug_assert!((tiles.len() as u32) * Self::BITS_PER_TILE < u32::BITS);
        Self {
            state: tiles
                .iter()
                .rev()
                .fold(0u32, |id, &t| (id << Self::BITS_PER_TILE) + t as u32),
            zpos: tiles
                .iter()
                .position(|&x| x == 0)
                .expect("Zero must be provided")
                .try_into()
                .expect("Index must fit in u32"),
        }
    }

    const fn index(r: u32, c: u32) -> u32 {
        debug_assert!(r < Self::HEIGHT);
        debug_assert!(c < Self::WIDTH);
        r * Self::WIDTH + c
    }

    const fn indices(pos: u32) -> (u32, u32) {
        debug_assert!(pos < Self::SIZE);
        (pos / Self::WIDTH, pos % Self::WIDTH)
    }

    const fn offset_from_zero(&self, dir: &Direction) -> Option<u32> {
        //! Cell to move zero to, if it's within the board.
        let (r, c) = Self::indices(self.zpos);
        let (dr, dc) = dir.as_offset();
        let (new_r, r_overflow) = r.overflowing_add_signed(dr);
        let (new_c, c_overflow) = c.overflowing_add_signed(dc);
        if !r_overflow && !c_overflow && new_r < Self::HEIGHT && new_c < Self::WIDTH {
            Some(Self::index(new_r, new_c))
        } else {
            None
        }
    }

    const fn move_zero_to(&self, nonzero_pos: u32) -> Self {
        //! Swap a zero with a value at `nonzero_pos`.
        //!
        //! `nonzero_pos` must be 4-adjacent to zpos.
        let zero_bits = self.zpos * Self::BITS_PER_TILE;
        let nonzero_bits = nonzero_pos * Self::BITS_PER_TILE;
        let nonzero = (self.state >> nonzero_bits) & Self::TILE_MASK;
        Self {
            state: self.state + (nonzero << zero_bits) - (nonzero << nonzero_bits),
            zpos: nonzero_pos,
        }
    }

    pub fn move_zero(&self, dir: &Direction) -> Option<Self> {
        self.offset_from_zero(dir).map(|pos| self.move_zero_to(pos))
    }
}

/// This is slightly faster than derive(Hash) since we don't need to consider zpos
impl std::hash::Hash for Board {
    fn hash<H: std::hash::Hasher>(&self, hasher_state: &mut H) {
        let Board { state, .. } = self;
        state.hash(hasher_state)
    }
}

/// This is slightly faster than derive(PartialEq) since we don't need to consider zpos
impl PartialEq for Board {
    fn eq(&self, other: &Self) -> bool {
        self.state.eq(&other.state)
    }
}

Let's test and benchmark!

#![feature(test)]
mod op; // `op.rs` contains your original code

// [snip] previous code here

fn main() {
}

#[cfg(test)]
mod bench {
    extern crate test;
    use super::{op::OPSolution, Solution};
    use test::{black_box, Bencher};

    #[bench]
    fn my_close(b: &mut Bencher) {
        let input = vec![vec![1, 2, 3], vec![4, 0, 5]];
        b.iter(|| black_box(Solution::sliding_puzzle(black_box(input.clone()))))
    }

    #[bench]
    fn op_close(b: &mut Bencher) {
        let input = vec![vec![1, 2, 3], vec![4, 0, 5]];
        b.iter(|| black_box(OPSolution::sliding_puzzle(black_box(input.clone()))))
    }

    #[bench]
    fn my_medium(b: &mut Bencher) {
        let input = vec![vec![4, 1, 2], vec![5, 0, 3]];
        b.iter(|| black_box(Solution::sliding_puzzle(black_box(input.clone()))))
    }

    #[bench]
    fn op_medium(b: &mut Bencher) {
        let input = vec![vec![4, 1, 2], vec![5, 0, 3]];
        b.iter(|| black_box(OPSolution::sliding_puzzle(black_box(input.clone()))))
    }

    #[bench]
    fn my_impossible(b: &mut Bencher) {
        let input = vec![vec![1, 2, 3], vec![5, 4, 0]];
        b.iter(|| black_box(Solution::sliding_puzzle(black_box(input.clone()))))
    }

    #[bench]
    fn op_impossible(b: &mut Bencher) {
        let input = vec![vec![1, 2, 3], vec![5, 4, 0]];
        b.iter(|| black_box(OPSolution::sliding_puzzle(black_box(input.clone()))))
    }
}

#[cfg(test)]
mod test {
    use super::{Board, Direction, Solution};

    #[test]
    fn test_move() {
        let start = Board::from_vec(&vec![1, 2, 3, 4, 0, 5]);
        let end = Board::from_vec(&vec![1, 2, 3, 4, 5, 0]);

        assert_eq!(start.move_zero(&Direction::Right), Some(end));
    }

    #[test]
    fn test_possible() {
        assert_eq!(
            Solution::sliding_puzzle(vec![vec![1, 2, 3], vec![4, 5, 0]]),
            0
        );
        assert_eq!(
            Solution::sliding_puzzle(vec![vec![1, 2, 3], vec![4, 0, 5]]),
            1
        );
        assert_eq!(
            Solution::sliding_puzzle(vec![vec![4, 1, 2], vec![5, 0, 3]]),
            5
        );
    }

    #[test]
    fn test_impossible() {
        assert_eq!(
            Solution::sliding_puzzle(vec![vec![1, 2, 3], vec![5, 4, 0]]),
            -1
        );
    }
}

$ cargo +nightly bench
   Compiling temp v0.1.0 (/tmp/temp)
    Finished `bench` profile [optimized] target(s) in 0.69s
     Running unittests src/main.rs (target/release/deps/temp-752be94e56df9355)

running 9 tests
test test::test_impossible ... ignored
test test::test_move ... ignored
test test::test_possible ... ignored
test bench::my_close      ... bench:         194.39 ns/iter (+/- 432.90)
test bench::my_impossible ... bench:      17,357.29 ns/iter (+/- 1,822.21)
test bench::my_medium     ... bench:         884.46 ns/iter (+/- 234.10)
test bench::op_close      ... bench:      40,205.81 ns/iter (+/- 7,938.55)
test bench::op_impossible ... bench:      40,220.88 ns/iter (+/- 6,439.14)
test bench::op_medium     ... bench:      39,983.23 ns/iter (+/- 21,219.03)

test result: ok. 0 passed; 0 failed; 3 ignored; 6 measured; 0 filtered out; finished in 11.47s

$ cargo +nightly test
   Compiling temp v0.1.0 (/tmp/temp)
    Finished `test` profile [unoptimized + debuginfo] target(s) in 0.19s
     Running unittests src/main.rs (target/debug/deps/temp-a77ab464bf664b98)

running 9 tests
test bench::my_close ... ok
test bench::my_medium ... ok
test bench::my_impossible ... ok
test test::test_impossible ... ok
test test::test_move ... ok
test test::test_possible ... ok
test bench::op_close ... ok
test bench::op_medium ... ok
test bench::op_impossible ... ok

test result: ok. 9 passed; 0 failed; 0 ignored; 0 measured; 0 filtered out; finished in 0.01s