Multiple conditional statements on list comprehension

Viewed 73

So this is my code and I want to know if I can use list comprehension to execute the same operation (count the clusters within rows and output a list of length df.shape[0]). There are at least two rows for the same cluster number, but it can be more and they cycles. I tried but couldn't figure it out. Any suggestions?

My code:

    import pandas as pd

    cluster_global = 0
    cluster_relativo = 0
    cluster_index = []
    for index, row in df.iterrows():
        if row['cluster'] == cluster_relativo:
            cluster_index.append(cluster_global)
        elif row['cluster'] == (cluster_relativo + 1):
            cluster_global += 1
            cluster_relativo += 1
            cluster_index.append(cluster_global)
        elif row['cluster'] == 0:
            cluster_global += 1
            cluster_relativo = 0
            cluster_index.append(cluster_global)

The DataFrame looks like

index cluster
0 0
1 0
2 1
3 1
4 1
5 2
6 2
7 0
8 0
... ...
n m<40
1 Answers

Do you want this?

from itertools import groupby

result = [0 if index == 0 and key == 0
          else index
          for index, (key, group) in enumerate(groupby(my_values))
          for _ in group
          ]

print(result)

Replace my_values in the list comprehension via - df['cluster'].values. to test

Related