Search code examples
pythonalgorithmmarkov-chainsmarkov

How to calculate the probability mass function of a random variable modulo N, where N is a prime number?


I'm trying to solve the following math problem:

A knight in standard international chess is sitting on a board as follows

0  1  2  3
4  5  6  7
8  9  10 11
12 13 14 15

The knight starts on square "0" and makes jumps to other squares according to the allowable moves in chess (so that at each space, it has between two to four valid moves). The knight chooses amongst the allowable moves at each jump uniformly at random and keeps track of the running sum S of keys on which it lands.

a. After T = 16 moves, what is the mean of the quantity S modulo 13?

b. What is the standard deviation?

c. After T = 512 moves, what is the mean of the quantity S modulo 311?

d. What is the standard deviation?

e. After T = 16 moves, what is the probability that the sum is divisible by 5, given that it is divisible by 13?

f. After T = 512 moves, what is the probability that the sum is divisible by 7, given that it is divisible by 43?

So far, I've written a program which calculates the probability mass function (pmf) of S:

from itertools import chain, product
import numpy as np
import pytest


def index_to_grid(index):
    return index // 4, index % 4

def grid_to_index(i, j):
    return 4*i + j

def in_board(i, j):
    return (0 <= i < 4) and (0 <= j < 4)

def available_moves(index):
    pos = np.array(index_to_grid(index))
    knight_hops = [np.array(hop) for hop in chain(product([-2, 2], [-1, 1]), product([-1, 1], [-2, 2]))]
    return set(grid_to_index(*newpos) for newpos in pos + knight_hops if in_board(*newpos))

def transition_matrix():
    T = np.zeros((16, 16))
    for i in range(16):
        js = available_moves(i)
        for j in js:
            T[i, j] = 1/len(js)
    return T

def calculate_S(N):
    '''Calculate the matrix S(i, n) of the expected value of S given initial state i after n transitions'''
    T = transition_matrix()
    S = np.zeros((16, N+1))
    for i in range(16):
        S[i, 0] = i

    # Use a bottom-up dynamic programming approach
    for n in range(1, N+1):
        for i in range(16):
            S[i, n] = sum(T[i, j] * (i + S[j, n-1]) for j in range(16))
    return S

Here are a few unit tests I've used to check my results so far:

def test_available_moves():
    assert available_moves(0) == {6, 9}
    assert available_moves(1) == {8, 10, 7}
    assert available_moves(10) == {4, 1, 12, 3}

def test_transition_matrix():
    T = transition_matrix()
    assert T[0, 6] == T[0, 9] == 1/2
    assert all(T[0, j] == 0 for j in set(range(16)) - {6, 9})
    assert T[1, 8] == T[1, 10] == T[1, 7] == 1/3
    assert all(T[1, j] == 0 for j in set(range(16)) - {8, 10, 7})
    assert T[10, 4] == T[10, 1] == T[10, 12] == T[10, 3] == 1/4
    assert all(T[10, j] == 0 for j in set(range(16)) - {4, 1, 12, 3})

def test_calculate_S():
    S = calculate_S(2)
    assert S[15, 1] == 15 + 1/2 * 6 + 1/2 * 9
    assert S[4, 1] == 4 + 1/3 * 2 + 1/3 * 10 + 1/3 * 13
    assert S[15, 2] == 15 + 1/2 * 9 + 1/2 * (1/4 * 0 + 1/4 * 2 + 1/4 * 7 + 1/4 * 15) \
                          + 1/2 * 6 + 1/2 * (1/4 * 0 + 1/4 * 8 + 1/4 * 13 + 1/4 * 15)


if __name__ == "__main__":
    pytest.main([__file__, "-s"])

So for example, to calculate the expected value of S itself after T = 16, I would evaluate calculate_S()[0, 16].

The problem is that I am having trouble generalizing this to the expected value of S % 13 (S modulo 13). Given that 13 (and all its 'equivalents' in subsequent questions) are all prime numbers, I suspect there is a key observation to be made using the 'primeness', but so far I haven't figured out what. Any ideas?


Solution

  • The trick is to use dynamic programming, and do all calculations mod some number. For each step you need the probability of it being at each square, with some sum mod some number.

    For example for problem f you need to do your sum calculations mod 7*43 = 301. So for each step you need the probabilities of being in all of the 16*301 = 4816 possible combinations of position and running sum mod 301.

    This makes your needed transition matrix much bigger.