I'm counting the probability of staying on an N by N chessboard after K steps, given that the starting position is r, c. Here's my code:
def knightProbability(self, N, K, r, c):
# dfs.
answer = 0
dp = [[[-1] * N]*N] * (K+1)
# dfs with memoization.
self.dic = {}
return self.dfs(N, K, r, c, dp) # Starting at r, c, what's the probability of staying on board after k moves.
def dfs(self, N, k, r, c, dp):
offsets = [(-1, -2), (-1, 2), (-2, -1), (-2, 1), (1, -2), (1, 2), (2, -1), (2, 1)]
answer = 0.0
if k == 0:
return 1.0
'''if dp[k][r][c] >= 0:
return dp[k][r][c]'''
if str(k) + ":" + str(r) + ":" + str(c) in self.dic:
return self.dic[str(k) + ":" + str(r) + ":" + str(c)]
for offset in offsets:
new_r = r + offset[0]
new_c = c + offset[1]
if new_r > -1 and new_r < N and new_c > -1 and new_c < N:
answer += self.dfs(N, k-1, new_r, new_c, dp) * 1.0 / 8 # with k-1 steps, what's the probability?
'''dp[k][r][c] = answer'''
self.dic[str(k) + ":" + str(r) + ":" + str(c)] = answer
return answer
It's a simple dfs. When I use a dictionary to store and retrieve visited states (k, r, c) like I'm doing now, it works. However, if I use a k by N by N array to store and retrive them, I'm not getting the correct answers. Any ideas why? The commented part is the storing/ retrieving from the array parts I'm talking about, which should be doing the same thing as the current code.