knights_probability - OS

knights_probability.py

class Solution:
    def knightProbability(self, N: int, K: int, r: int, c: int) -> float:
        moves = [(-2, 1), (-2, -1), (-1, 2), (-1, -2), (1, 2), (1, -2), (2, 1), (2, -1)]
        dp0 = [[1] * N for _ in range(N)]
        to_visit = {(r, c)}
        count = 0
        for _ in range(K):
            dp1 = [[0] * N for _ in range(N)]
            new_to_visit = set()
            for (r, c) in to_visit:
                for (dr, dc) in moves:
                    nr, nc = r + dr, c + dc
                    if 0 <= nr < N and 0 <= nc < N:
                        dp1[nr][nc] += dp0[r][c]
                        new_to_visit.add((nr, nc))
            to_visit = new_to_visit
            dp0 = dp1
        count = sum(map(sum, dp0))
        return count / float(8 ** K)