import numpy as np
import matplotlib.pyplot as plt
import math

Nmax=100#maximální počet kroků
p=0.5 #pravděpodobnost, že udělám krok vpravo
N_repeat=1000 #počet opakování simulací
N=np.arange(0,Nmax+1,1) #pole celkových počtů kroků
Psim=np.zeros(np.size(N)) #pravděpodobnost, že se vrátím do počátku - odhad ze simulace
for i in N:
    data=np.random.binomial(N[i],p,N_repeat) #simulace počtu kroků vpravo pomocí binomického rozdělení
    for k in data:   #počítám kolikrát se vrátím do počátku
        if (2*k-N[i])==0:
            Psim[i]+=1
Psim=Psim/N_repeat    

Pteor=np.zeros(np.size(N))
for i in N:
    Pteor[i]=math.comb(int(N[i]),int(N[i]/2))*(1/2)**N[i] #pravděpodobnost, že se vrátím do počátku - teoretická hodnota vypočtená z binomického rozdělení

plt.step(N,Psim,label='simulace')
plt.plot(N,Pteor,c='red',label='teorie')
plt.xlim(0,Nmax)
plt.ylim(0,1)
plt.xlabel('počet kroků')
plt.ylabel('pravděpodobnost')
plt.legend()
plt.show()