print_word_search_expectation.py

import sympy

def make_id(size):
    return sympy.Matrix([[int(i==j) for j in xrange(size)] for i in xrange(size)])

def add_absorb(k, m):
    size = len(m)
    assert size == len(m[0])
    return [row + [0] for row in m] + [ [0]* (size-1) + [k,1] ]

def get_matrix_before_success_expectation(m, w, k):
    I = make_id(w)
    X = m / k
    tmp = (I-X).inv()
    tmp.simplify()
    R = (tmp*tmp + tmp * (w-1)) / (k**(w))
    R.simplify()
    return R

def f(w):
    k = sympy.var("k")
    m = [[k-1] * w] + [[int(j == i) for j in xrange(w)] for i in xrange(w-1)]
    m = sympy.Matrix(m)
    m = sympy.Matrix(m)

    m2 = [[k-2] * w] + [[1] * w] + [[int(j == i+1) for j in xrange(w)] for i in xrange(w-2)]
    m2[0][0] += 1
    m2 = sympy.Matrix(m2)

    print
    print 'w =', w
    a = get_matrix_before_success_expectation(m, w, k)[0,0]
    tmp = get_matrix_before_success_expectation(m2, w, k)
    b = sympy.simplify(sum(tmp[i,0] for i in xrange(w)))

    print a
    print b


def main():
    for w in xrange(2, 10):
        f(w)

main()

#output:

# w = 2
# k*(k + 1)
# k**2

# w = 3
# k*(k**2 + k + 1)
# k**3

# w = 4
# k*(k**3 + k**2 + k + 1)
# k**4

# w = 5
# k*(k**4 + k**3 + k**2 + k + 1)
# k**5

# w = 6
# k*(k**5 + k**4 + k**3 + k**2 + k + 1)
# k**6

# w = 7
# k*(k**6 + k**5 + k**4 + k**3 + k**2 + k + 1)
# k**7

# w = 8
# k*(k**7 + k**6 + k**5 + k**4 + k**3 + k**2 + k + 1)
# k**8

# w = 9
# k*(k**8 + k**7 + k**6 + k**5 + k**4 + k**3 + k**2 + k + 1)
# k**9