"""
Example from lecture on Sept 23. Small extension of the
interest problem, now with varying interest rate. p is
now an array, being 5% for the first three years 
and then 3%. 
"""


import numpy as np
import matplotlib.pyplot as plt

x0 = 100                      # initial amount
N = 40                         # number of years

years = range(N+1)
x = np.zeros(N+1)
p = np.ones(N)
p[:3] *= 5
p[3:] *=3
print(p)

x[0] = x0
for n in years[1:]:
    x[n] = x[n-1] + (p[n-1]/100.0)*x[n-1]
plt.plot(years, x, 'ro')
plt.xlabel('Years')
plt.ylabel('Amount')

plt.show()