In [1]:
def F(n,p):
    """
    n 本の傘を持っているとする。
    n + 1 本目の置き傘ができるまでの移動回数の期待値を求める。
    """
    import numpy as np
    q = 1 - p
    A = np.eye(n+1)
    A[np.arange(n+1), np.arange(n+1)[::-1]] -= q
    A[np.arange(1,n+1), np.arange(1,n+1)[::-1]] -= p
    x = np.linalg.solve(A, np.ones(n+1))
    return x[-1]
In [2]:
F(10 ** 5, 0.3)
---------------------------------------------------------------------------
MemoryError                               Traceback (most recent call last)
<ipython-input-2-a8d44d3f5129> in <module>
----> 1 F(10 ** 5, 0.3)

<ipython-input-1-f61f0188c547> in F(n, p)
      6     import numpy as np
      7     q = 1 - p
----> 8     A = np.eye(n+1)
      9     A[np.arange(n+1), np.arange(n+1)[::-1]] -= q
     10     A[np.arange(1,n+1), np.arange(1,n+1)[::-1]] -= p

C:\ProgramData\Anaconda3\lib\site-packages\numpy\lib\twodim_base.py in eye(N, M, k, dtype, order)
    199     if M is None:
    200         M = N
--> 201     m = zeros((N, M), dtype=dtype, order=order)
    202     if k >= M:
    203         return m

MemoryError: 
In [3]:
def F_sparse(n,p):
    """
    n 本の傘を持っているとする。
    n + 1 本目の置き傘ができるまでの移動回数の期待値を求める。
    """
    from scipy.sparse import identity, csr_matrix
    from scipy.sparse.linalg import spsolve
    import numpy as np
    q = 1 - p
    A = identity(n+1)
    A -= csr_matrix(([q] * (n+1), (np.arange(n+1), np.arange(n+1)[::-1])), (n+1, n+1))
    A -= csr_matrix(([p] * n, (np.arange(1,n+1), np.arange(1,n+1)[::-1])), (n+1, n+1))
    x = spsolve(A, np.ones(n+1))
    return x[-1]
In [4]:
%time F_sparse(10 ** 5, 0.3)
Wall time: 310 ms
Out[4]:
476193.60241394595