Search code examples
pythondataframematplotlibpylot

Change hardcode to more flexible using python for loop


I'm writing a code about plotting. I write with hardcode ways, so my code is not flexible enough.

I know can using for loop to solve the hardcode problems. But my Python ability not strong enough.

Here's my code.

df1 = df[df.cluster==0]
df2 = df[df.cluster==1]
df3 = df[df.cluster==2]

plt.scatter(df1.Age,df1['Income($)'],color='green')
plt.scatter(df2.Age,df2['Income($)'],color='red')
plt.scatter(df3.Age,df3['Income($)'],color='black')

On this case's there's have 3 cluster. If cluster = 4, then need write more. df4 = ...

Can i write a for loop, such like this

n = number of cluster
for i in range(n):
    df(random) = df[df.cluster==i]
for j in range(n):
    plt.scatter(df(n).Age,df(n)['Income($)'],color='RANDOM')

My question is write a code only a few line not using hardcode ways.


Solution

  • If you are looking for a simple solution, this might be it. (I have reused your code sample)

    n = num_of_clusters
    my_colors = ['green', 'red', 'black', ...]
    for i in range(n):
        df_i = df[df.cluster == i]
        plt.scatter(df_i.Age, df_i['Income($)'], color=my_colors[i])