import numpy as np
import matplotlib.pyplot as plt

Nmax=10000 #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ů
mean_distance=np.zeros(np.size(N)) #průměrná vzdálenost od středu 
mean_distance_abs=np.zeros(np.size(N)) #průměrná vzdálenost od středu 
for i in N:
    data=np.random.binomial(N[i],p,N_repeat) #simulace N_repeat hodnot z binomick0ho rozdělení B(N,p)
    mean_distance[i]=np.sqrt(np.mean((2*data-N[i])**2)) #výpočet průměrné hodnoty vzdálenosti od středu jako sqrt(<(2k-N)**2>)
    mean_distance_abs[i]=np.mean(abs((2*data-N[i]))) #výpočet průměrné hodnoty vzdálenosti od středu jako |2k-N|

plt.plot(N,mean_distance,label='simulace: odmocnina z kvadrátu souřadnice')
plt.plot(N,mean_distance_abs,c='green', label='simulace: absolutní hodnota')
plt.plot(N,np.sqrt(N),c='red',label='teorie: odmocnina z kvadrátu souřadnice')
plt.plot(N,np.sqrt(2*N/np.pi),c='violet',label='teorie: absolutní hodnota')
plt.xlim(0,Nmax)
plt.ylim(0,)
plt.xlabel('počet kroků')
plt.ylabel('vzdálenost')
plt.legend()
plt.show()