diff --git a/k-bandits.py b/k-bandits.py
index 8a751ce..96a428e 100644
--- a/k-bandits.py
+++ b/k-bandits.py
@@ -5,23 +5,25 @@ from typing import Callable
 float_formatter = "{:.3f}".format
 np.set_printoptions(formatter={'float_kind': float_formatter})
 
-K = 10
-
 def bandit(a: int, bias: np.array):
 	zn = np.random.normal() + bias[a]
 	return zn
 
 
-def test(_episode_len: int, _qu: np.array, _nu: np.array, _bandit: Callable, _param: tuple[float, float]) -> np.array:
+def test(_k_arms: int, _episode_len: int, _bandit: Callable, _param: tuple[float, float]) -> np.array:
 	_epsilon = _param[0]		# Anti-greediness (ability to explore)
 	_rho = _param[1]		# reduce epsilon with age
 	_r_mean = []
 
+	# Init Q and N
+	_qu = np.zeros(_k_arms)
+	_nu = np.zeros(_k_arms)
+
 	# Calc z in advance
 	z = np.random.uniform(size=_episode_len)
 
 	# Calc a_expl in advance
-	a_expl = np.random.randint(low=0, high=K, size=_episode_len)
+	a_expl = np.random.randint(low=0, high=_k_arms, size=_episode_len)
 
 	for j in range(0, _episode_len):
 		a = a_expl[j]
@@ -47,6 +49,8 @@ def test(_episode_len: int, _qu: np.array, _nu: np.array, _bandit: Callable, _pa
 if __name__ == '__main__':
 
 	# Init parameters
+	k_arms = 10
+	highest_reward = 1.5
 	num_realisations = 2000
 	episode_len = 1000
 
@@ -58,23 +62,20 @@ if __name__ == '__main__':
 		for realization in range(0, num_realisations):
 			# init bandits with different biases for shifting reward probability
 			# -> expected reward q*(a)
-			ql_star = np.random.uniform(-1.5, +1.5, K)
+			ql_star = np.random.uniform(-highest_reward, +highest_reward, k_arms)
 
-			# Init Q and N
-			qu = np.zeros(K)
-			nu = np.zeros(K)
-			r_mean += test(episode_len, qu, nu, _bandit=lambda a: bandit(a, ql_star), _param=param)
+			r_mean += test(_k_arms=k_arms, _episode_len=episode_len, _bandit=lambda a: bandit(a, ql_star), _param=param)
 
 		pl.plot(r_mean/num_realisations)
 		legend.append(f"Param {param}")
 
 	pl.grid()
-	pl.title(f"Norm. E(R), num. realizations: Z={num_realisations}, num. bandits: K={K}")
+	pl.title(f"Norm. E(R), num. realizations: Z={num_realisations}, num. bandit arms: K={k_arms}")
 	pl.xlabel("Episode")
 	pl.ylabel("E(R)/Z")
 	pl.legend(legend)
 	ax = pl.gca()
 	ax.set_xlim([-episode_len/10, episode_len])
-	ax.set_ylim([0, 1.5])
+	ax.set_ylim([0, highest_reward])
 
 	pl.show()