plot_cos(kx).py

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

n = int(input())

##generation
x = np.arange(0, 2*pi+0.1, 0.1)
array = [np.cos(i*x).tolist() for i in range(1, n+1, 1) ]
x = x.tolist()

## question one
graph = plt.figure()
ax = graph.add_subplot(111)
for i in range(n) :
    ax.plot(x, array[i], label = "k = " + str(i+1))
plt.legend()


## question two
for i in range(n) :
    graph = plt.figure()
    ax = graph.add_subplot(111)
    ax.plot(x, array[i], label = "k = " + str(i+1))
    plt.legend()
plt.show()