Python nested loop over pandas dataframe

Question

I have a pandas data frame of 30 000 rows that looks like this:

ID     year    month    var1-var300    test   

1111   2017    7        ...            1      
1111   2017    9        ...            0      
2222   2017    6        ...            1      
2222   2017    6        ...            0      
2222   2016    6        ...            0      
3333   2017    3        ...            1      
3333   2017    3        ...            0     
3333   2015    8        ...            0      
...

Here is what I want to do for each row: if test=1, I would like to extract the variables "ID year month", loop over the entire data frame and if this combination of variables is found in any other row, assign 1 to a new variable 'check'. The final data frame should look like this:

ID     year    month    var1-var300    test   check
1111   2017    7        ...            1      0
1111   2017    9        ...            0      0
2222   2017    6        ...            1      1
2222   2017    6        ...            0      0
2222   2016    6        ...            0      0
3333   2017    3        ...            1      1
3333   2017    3        ...            0      0
3333   2015    8        ...            0      0
...

Here is some kind of pseudo-code I have imagined:

for line in df:
    if line['test']=1:
        I=line['ID']
        Y=line['year']
        MO=line['month']
        for row in df:
            if row['ID']=I & row['year']=Y & row['month']=MO:
                line['check']=1
                break

Any idea how to do a similar code that works in Python?

you should not use assignment operator "=" there, instead use comparison operator "==". — Prakriti Shaurya
– Prakriti Shaurya, Commented Nov 25, 2020 at 10:05

Nils Werner · Accepted Answer · 2020-11-25 09:51:52Z

1

You should be able to invert your logic:

Group by ID and year
Do your check inside each group

def func(group):
    if len(group) > 1:
        group.loc[group['test'] == 1, 'check'] = 1
    return group

df = df.groupby(['ID', 'year']).apply(func)

answered Nov 25, 2020 at 9:51

Nils Werner

37.2k7 gold badges85 silver badges108 bronze badges

Sign up to request clarification or add additional context in comments.

Comments

Xu Qiushi · Accepted Answer · 2020-11-25 10:04:00Z

1

I think you can just use a transform to count group. Then you can get the result. Just two lines.

This is my solution.

Create Test Data:

import pandas as pd
ID = [1111, 1111, 2222, 2222, 2222, 3333, 3333, 3333]
year = [2017, 2017, 2017, 2017, 2016, 2017, 2017, 2015]
month = [7, 9, 6, 6, 6, 3, 3, 8]
test = [1, 0, 1, 0, 0, 1, 0, 0]
df = pd.DataFrame({
    "ID": ID,
    "year": year,
    "month": month,
    "test": test
})

Get the result:

df.loc[:, "group_count"] = df.groupby(["ID", "year", "month"]).transform("count").values
df.loc[:, "check"] = ((df["test"]>0) & (df["group_count"] > 1)).astype(int)

answered Nov 25, 2020 at 10:04

Xu Qiushi

1,1711 gold badge7 silver badges11 bronze badges

3 Comments

Falco Over a year ago

Interesting approach...tank you! :-)

Falco Over a year ago

I tried the first row in your solution but I got an error. Here is a version that works: df["group_count"] = df.groupby(["ID", "year", "month"])["ID"].transform("count")

Xu Qiushi Over a year ago

@Falco I did not got error on my local machine. I guess may be it is because version. My python 3.7, pandas 0.25.1.

a_gdevr · Accepted Answer · 2020-11-25 10:02:16Z

0

So you want a single column where it is indicated if the corresponding row's ID, year and month correspond to ID, year and month of a row with test == 1?

You iterate with iterrows():

to_check = []
for index, row in df.iterrows():
    if row['test']==1: # in your pseudocode, you use single =; that's for assigning variables
        to_ckeck.append([row['ID'], row['year'], row['month']])

check = []
for index, row in df.iterrows():
      if [row['ID'], row['year'], row['month']] in to_check:
          check.append(1)
      else:
          check.append(0)
df["check"] = check

edited Nov 25, 2020 at 10:02

answered Nov 25, 2020 at 9:52

a_gdevr

1031 silver badge7 bronze badges

Comments

Prakriti Shaurya · Accepted Answer · 2020-11-25 10:12:10Z

0

You can make some changes as below and try:

for line in df:
    if line['test']==1:
        I=line['ID']
        Y=line['year']
        MO=line['month']
        for row in df:
            if row['ID']==I and row['year']==Y and row['month']==MO:
                line['check']=1
                break

answered Nov 25, 2020 at 10:12

Prakriti Shaurya

1951 gold badge3 silver badges15 bronze badges

Collectives™ on Stack Overflow

Python nested loop over pandas dataframe

4 Answers 4

Comments

3 Comments

Comments

Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

4 Answers 4

Comments

3 Comments

Comments

Comments

Your Answer

Sign up or log in

Post as a guest

Related