Instantly share code, notes, and snippets.

# RuolinZheng08/backtracking_template.py

Last active June 6, 2024 22:47
Show Gist options
• Save RuolinZheng08/cdd880ee748e27ed28e0be3916f56fa6 to your computer and use it in GitHub Desktop.
[Algo] Backtracking Template & N-Queens Solution
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters. Learn more about bidirectional Unicode characters
 def is_valid_state(state): # check if it is a valid solution return True def get_candidates(state): return [] def search(state, solutions): if is_valid_state(state): solutions.append(state.copy()) # return for candidate in get_candidates(state): state.add(candidate) search(state, solutions) state.remove(candidate) def solve(): solutions = [] state = set() search(state, solutions) return solutions
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters. Learn more about bidirectional Unicode characters
 class Solution: """ example on the left: [1, 3, 0, 2] example on the right: [2, 0, 3, 1] """ def solveNQueens(self, n: int) -> List[List[str]]: solutions = [] state = [] self.search(state, solutions, n) return solutions def is_valid_state(self, state, n): # check if it is a valid solution return len(state) == n def get_candidates(self, state, n): if not state: return range(n) # find the next position in the state to populate position = len(state) candidates = set(range(n)) # prune down candidates that place the queen into attacks for row, col in enumerate(state): # discard the column index if it's occupied by a queen candidates.discard(col) dist = position - row # discard diagonals candidates.discard(col + dist) candidates.discard(col - dist) return candidates def search(self, state, solutions, n): if self.is_valid_state(state, n): state_string = self.state_to_string(state, n) solutions.append(state_string) return for candidate in self.get_candidates(state, n): # recurse state.append(candidate) self.search(state, solutions, n) state.pop() def state_to_string(self, state, n): # ex. [1, 3, 0, 2] # output: [".Q..","...Q","Q...","..Q."] ret = [] for i in state: string = '.' * i + 'Q' + '.' * (n - i - 1) ret.append(string) return ret
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters. Learn more about bidirectional Unicode characters
 class Solution: from itertools import product SHAPE = 9 GRID = 3 EMPTY = '.' DIGITS = set([str(num) for num in range(1, SHAPE + 1)]) def solveSudoku(self, board: List[List[str]]) -> None: """ Do not return anything, modify board in-place instead. """ self.search(board) def is_valid_state(self, board): # check if it is a valid solution # validate all the rows for row in self.get_rows(board): if not set(row) == self.DIGITS: return False # validate columns for col in self.get_cols(board): if not set(col) == self.DIGITS: return False # validate sub-boxes for grid in self.get_grids(board): if not set(grid) == self.DIGITS: return False return True def get_candidates(self, board, row, col): used_digits = set() # remove digits used by the same row used_digits.update(self.get_kth_row(board, row)) # remove digits used by the same column used_digits.update(self.get_kth_col(board, col)) # remove digits used by the 3x3 sub-box used_digits.update(self.get_grid_at_row_col(board, row, col)) used_digits -= set([self.EMPTY]) candidates = self.DIGITS - used_digits return candidates def search(self, board): if self.is_valid_state(board): return True # found solution # find the next empty spot and take a guess for row_idx, row in enumerate(board): for col_idx, elm in enumerate(row): if elm == self.EMPTY: # find candidates to construct the next state for candidate in self.get_candidates(board, row_idx, col_idx): board[row_idx][col_idx] = candidate # recurse on the modified board is_solved = self.search(board) if is_solved: return True else: # undo the wrong guess and start anew board[row_idx][col_idx] = self.EMPTY # exhausted all candidates # but none solves the problem return False # no empty spot return True # helper functions for retrieving rows, cols, and grids def get_kth_row(self, board, k): return board[k] def get_rows(self, board): for i in range(self.SHAPE): yield board[i] def get_kth_col(self, board, k): return [ board[row][k] for row in range(self.SHAPE) ] def get_cols(self, board): for col in range(self.SHAPE): ret = [ board[row][col] for row in range(self.SHAPE) ] yield ret def get_grid_at_row_col(self, board, row, col): row = row // self.GRID * self.GRID col = col // self.GRID * self.GRID return [ board[r][c] for r, c in product(range(row, row + self.GRID), range(col, col + self.GRID)) ] def get_grids(self, board): for row in range(0, self.SHAPE, self.GRID): for col in range(0, self.SHAPE, self.GRID): grid = [ board[r][c] for r, c in product(range(row, row + self.GRID), range(col, col + self.GRID)) ] yield grid

Awesome.

### kyzmitch commented May 25, 2022

thanks for the template