Plotting line graph for 3 variables in python

Viewed 7961
Country Month   GDPA
USA     January  5
USA     February 10
USA     March    5
Canada  January  7
Canada  February 4
Canada  March    10
Mexico  January  9
Mexico  February 17
Mexico  March    8

Above is my data frame.

I want to make a line graph of the df. Where the x axis is the variable Month. The y axis is GDPA. And there are three separate lines for each country. Preferably using matplotlib only.

plt.plot(df[‘Country’], df[‘GDPA’], df[‘Month’]

My code above does not produce the expected graph because the lines are connected and it doesn’t even take the GDPA values into account.

I want it so that the x axis is: January February March
y axis is GDPA values.
and there are 3 lines (for each country)

enter image description here

1 Answers

This may not be as elegant as possible, but a low-level way to do it is to basically just make a set of the countries (so you have no duplicates), then loop over that set and plot a line for each country.

import matplotlib.pyplot as plt
import pandas

d = {'Country': ['USA', 'USA', 'USA', 'Canada', 'Canada', 'Canada', 'Mexico', 'Mexico', 'Mexico'],
     'Month': ['January', 'February', 'March', 'January', 'February', 'March', 'January', 'February', 'March'],
     'GDPA': [5, 10, 5, 7, 4, 10, 9, 17, 8]}
df = pandas.DataFrame(data=d)
print(df)

country_set = set(df['Country'])

plt.figure()
for country in country_set:
     selected_data = df.loc[df['Country'] == country]
     plt.plot(selected_data['Month'], selected_data['GDPA'], label=country)
     
plt.legend()
plt.show()

Resulting plot here

Pure python and matplotlib!

Related