Search code examples
pythonlistmatplotlibzip

Pythonic way to extract a subset of lists to plot


I have this code which works fine but I am sure there is a better way to do it. I have some lists which I zip. Then based on a condition I make a subset of that zip list. Then I want to plot both . The subset is plotted in a different colour.

Code

import matplotlib
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt

a = [11,12,13,14,15,16,17,18,19,20]
b = [103,121,99,133,89,107,102,106,110,105]
ind = [11,14,18]
c = list(zip(a,b))

#output = [(11,103),(14,133),(18,106)]
print(':', c, '\n')



e = []
for x,y in zip(a,b):
    if x in ind:
        e.append((x,y))
        print(x, y)
print(e)
f = list(zip(*e))
x = f[0]
y = f[1]
print('\n', 'x :', x)
print('\n', 'f:',f)

# graph 
fig1=plt.figure(num=None, figsize=(14, 6), dpi=80, facecolor='w', edgecolor='k')
plt.plot( a, b, "o", color="blue", markersize=4, label = '0')
plt.plot( x, y, "o", color="red", markersize=4, label = '0')
plt.show()

Desired output

A more pythonic way of achieving the same result.


Solution

  • filtered_pairs = [e for e in zip(a, b) if e[0] in ind]
    x = [e[0] for e in filtered_pairs]
    y = [e[1] for e in filtered_pairs]
    

    List comprehensions are very pythonic.