Problem statement

https://leetcode.com/problems/knight-probability-in-chessboard/

Solution

Let us use dp[k][i][j] be probability to be on (i, j) cell on k moment of time. Then we can update all table in O(8N^2K): for each cell look at all its knight neighbors.

Complexity

Time complexity is O(8n^2K), space complexity is O(n^2K)

Code

class Solution:
    def knightProbability(self, N, K, r, c):
        dp = [[[0] * N for _ in range(N)] for _ in range(K+1)]
        dp[0][r][c] = 1
        neibs = [[1,2],[1,-2],[-1,2],[-1,-2],[2,1],[2,-1],[-2,1],[-2,-1]]
        for k,i,j in product(range(K),range(N),range(N)):
            for dx, dy in neibs:
                if 0 <= i+dx < N and 0 <= j+dy < N:
                    dp[k+1][i][j] += dp[k][i+dx][j+dy] / 8
        
        return sum(chain(*dp[-1]))

Remark

There is also O(N^6 * log K) time complexity solution, if we use N^2 x N^2 transition matrix (which is very sparse in fact) and look at our problem as Markov Chain.