Simple Sudoku Solver

sudoku_solver.rs

use colored::*;
use std::fmt;
use std::io::{self, BufRead};

#[derive(Clone)]
struct Cell {
  possibilities: String,
}

impl Default for Cell {
  fn default() -> Self {
    Cell {
      possibilities: "123456789".to_string(),
    }
  }
}

impl fmt::Debug for Cell {
  fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
    write!(f, "{}", self.possibilities)
  }
}

#[derive(Clone)]
struct Grid {
  cells: Vec<Cell>,
  pre_filled: Vec<usize>,
}

impl Grid {
  fn new() -> Self {
    Grid {
      cells: vec![Cell::default(); 81],
      pre_filled: Vec::<usize>::with_capacity(17),
    }
  }

  fn parse(&mut self, input: &str) -> bool {
    debug_assert_eq!(input.len(), 81);
    for (idx, ch) in input.char_indices() {
      if ch.is_ascii_digit() {
        if !self.fill(idx, ch) {
          return false;
        } else {
          self.pre_filled.push(idx);
        }
      }
    }
    true
  }

  fn row_col(idx: usize) -> (usize, usize) {
    (idx / 9, idx % 9)
  }

  fn idx(row: usize, col: usize) -> usize {
    row * 9 + col
  }

  fn sub(row: usize, col: usize, i: usize) -> usize {
    const START_IDX: [usize; 9] = [0, 3, 6, 27, 30, 33, 54, 57, 60];
    let block_idx = (row / 3) * 3 + (col / 3);
    let sub_start_idx = START_IDX[block_idx];
    sub_start_idx + ((i / 3) * 9) + (i % 3)
  }

  fn row_peers(cell: usize) -> Vec<usize> {
    let row = cell / 9;
    (0..9)
      .map(|i| Grid::idx(row, i))
      .filter(|&c| c != cell)
      .collect()
  }

  fn col_peers(cell: usize) -> Vec<usize> {
    let col = cell % 9;
    (0..9)
      .map(|i| Grid::idx(i, col))
      .filter(|&c| c != cell)
      .collect()
  }

  fn block_peers(cell: usize) -> Vec<usize> {
    let (row, col) = Grid::row_col(cell);
    (0..9)
      .map(|i| Grid::sub(row, col, i))
      .filter(|&c| c != cell)
      .collect()
  }

  fn peers(idx: usize) -> Vec<usize> {
    let mut row_peers = Grid::row_peers(idx);
    let mut col_peers = Grid::col_peers(idx);
    let mut peers = Grid::block_peers(idx);
    peers.append(&mut col_peers);
    peers.append(&mut row_peers);
    peers.sort_unstable();
    peers.dedup();
    peers
  }

  // returns [None, None] if there’re no occurances of |digit|
  // returns [Some(idx), None] if there’s only one occurance of it
  // returns [Some(idx), Some(idx)] if digit has more than one occurance
  fn digit_occurances(
    &self,
    indices: &Vec<usize>,
    digit: char,
  ) -> [Option<usize>; 2] {
    let mut peers_with_digit = [None; 2];
    let mut i = 0;
    for &idx in indices
      .iter()
      .filter(|&&k| self.cells[k].possibilities.find(digit).is_some())
      // we aren’t interested if |digit| occurs more than once
      .take(2)
    {
      peers_with_digit[i] = Some(idx);
      i += 1;
    }
    peers_with_digit
  }

  fn fill(&mut self, idx: usize, digit: char) -> bool {
    let remaining = self.cells[idx]
      .possibilities
      .replace(&digit.to_string(), "");
    remaining.chars().all(|c| self.eliminate(idx, c))
  }

  fn eliminate(&mut self, idx: usize, digit: char) -> bool {
    let cell = &mut self.cells[idx];
    if cell.possibilities.find(digit).is_none() {
      return true; // already eliminated
    }
    if cell.possibilities.len() == 1 {
      return false; // eliminated all possibilities
    }
    cell.possibilities = cell.possibilities.replace(&digit.to_string(), "");

    if cell.possibilities.len() == 1 {
      // this cell now has an assigned value, notify peers
      let final_digit = cell.possibilities.chars().next().unwrap();
      if !Grid::peers(idx)
        .iter()
        .all(|&p| self.eliminate(p, final_digit))
      {
        return false;
      }
    }

    // if |digit| has only one possible location amongst peers in units
    // fill it in that peer
    for peers in Grid::unit_iterator(idx) {
      match self.digit_occurances(&peers, digit) {
        // digit has no place to fill, contradiction!
        [None, None] => return false,
        // digit has just one possible cell for this unit, fill it
        [Some(peer_idx), None] => {
          return self.fill(peer_idx, digit);
        }
        _ => (),
      }
    }

    true
  }

  fn is_unit_solved(u: &str) -> bool {
    if u.len() != 9 {
      return false;
    }
    let mut counts = [0_u8; 9];
    for ch in u.chars() {
      let idx = ch.to_digit(10).unwrap() as usize - 1;
      counts[idx] += 1;
    }
    counts.iter().all(|&c| c == 1)
  }

  fn get_row(&self, idx: usize) -> String {
    let row_start = 9 * idx;
    self.cells[row_start..(row_start + 9)]
      .iter()
      .map(|c| c.possibilities.clone())
      .fold(String::new(), |acc, cell| acc + &cell)
  }

  fn get_col(&self, idx: usize) -> String {
    self.cells[idx..]
      .iter()
      .step_by(9)
      .map(|c| c.possibilities.clone())
      .fold(String::new(), |acc, cell| acc + &cell)
  }

  fn get_block(&self, idx: usize) -> String {
    let (row, col) = Grid::row_col(idx);
    let mut result = String::new();
    for i in 0..9 {
      let j = Grid::sub(row, col, i);
      result.push_str(&self.cells[j].possibilities);
    }
    result
  }

  fn is_solved(&self) -> bool {
    (0..9).all(|i| Grid::is_unit_solved(&self.get_row(i)))
      && (0..9).all(|i| Grid::is_unit_solved(&self.get_col(i)))
      && (0..9).all(|i| Grid::is_unit_solved(&self.get_block(i)))
  }

  fn get_next_unfilled(&self) -> Option<usize> {
    let mut min_idx = 81;
    let mut min_possibilities = 9;
    for (i, c) in self.cells.iter().enumerate() {
      if c.possibilities.len() < 2 {
        continue;
      }
      if c.possibilities.len() == 2 {
        return Some(i);
      }
      if min_possibilities > c.possibilities.len() {
        min_possibilities = c.possibilities.len();
        min_idx = i;
      }
    }
    if min_possibilities < 9 {
      return Some(min_idx);
    }
    None
  }

  fn unit_iterator(cell_idx: usize) -> CellUnitIterator {
    CellUnitIterator {
      unit_index: 0,
      cell_index: cell_idx,
    }
  }
}

struct CellUnitIterator {
  unit_index: usize,
  cell_index: usize,
}

impl Iterator for CellUnitIterator {
  type Item = Vec<usize>;
  fn next(&mut self) -> Option<Self::Item> {
    self.unit_index += 1;
    match self.unit_index - 1 {
      0 => Some(Grid::row_peers(self.cell_index)),
      1 => Some(Grid::col_peers(self.cell_index)),
      2 => Some(Grid::block_peers(self.cell_index)),
      _ => None,
    }
  }
}

impl fmt::Debug for Grid {
  fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
    let max_cell_len = self
      .cells
      .iter()
      .max_by_key(|&c| c.possibilities.len())
      .unwrap()
      .possibilities
      .len();
    for i in 0..9 {
      for j in 0..9 {
        write!(
          f,
          "{:^width$} ",
          format!("{}", self.cells[Grid::idx(i, j)].possibilities),
          width = max_cell_len
        )
        .unwrap();
        if (j + 1) % 3 == 0 && (j != 8) {
          write!(f, "| ").unwrap();
        }
      }
      writeln!(f, "").unwrap();
      if i == 2 || i == 5 {
        let n = (max_cell_len + 1) * 9 + 3;
        let separator: String = (0..n).map(|_| "-").collect();
        writeln!(f, "{}", separator).unwrap();
      }
    }
    Ok(())
  }
}

impl fmt::Display for Grid {
  fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
    let max_cell_len = self
      .cells
      .iter()
      .max_by_key(|&c| c.possibilities.len())
      .unwrap()
      .possibilities
      .len();
    for i in 0..9 {
      for j in 0..9 {
        let idx = Grid::idx(i, j);
        let printed = match self.pre_filled.contains(&idx) {
          true => self.cells[idx].possibilities.red(),
          false => self.cells[idx].possibilities.normal(),
        };
        write!(f, "{:^width$} ", printed, width = max_cell_len).unwrap();
        if (j + 1) % 3 == 0 && (j != 8) {
          write!(f, "| ").unwrap();
        }
      }
      writeln!(f, "").unwrap();
      if i == 2 || i == 5 {
        let n = (max_cell_len + 1) * 9 + 3;
        let separator: String = (0..n).map(|_| "-").collect();
        writeln!(f, "{}", separator).unwrap();
      }
    }
    Ok(())
  }
}

fn search(current_board: &Grid, final_board: &mut Grid) -> bool {
  if let Some(unfilled_idx) = current_board.get_next_unfilled() {
    let current_cell = &current_board.cells[unfilled_idx];
    for digit in current_cell.possibilities.chars() {
      let mut new_board = current_board.clone();
      if new_board.fill(unfilled_idx, digit) {
        if search(&new_board, final_board) {
          return true;
        }
      }
    }
  } else {
    // no more squares, we’re done!
    debug_assert!(current_board.is_solved());
    *final_board = current_board.clone();
    return true;
  }
  false
}

fn main() {
  let puzzles: Vec<String> =
    io::stdin().lock().lines().map(|l| l.unwrap()).collect();

  for puzzle in &puzzles {
    let mut g = Grid::new();
    g.parse(puzzle);
    println!("{:?}", g);
    let mut solved = Grid::new();
    if search(&g, &mut solved) {
      println!("{}", solved);
    } else {
      println!("Invalid board.");
    }
  }
}

Rust translation of Peter Norvig's sudoku solver in Python.

Of course, lots of optimisations are possible! This is just a pedagogical example 🤓