Identifying specific strings and checking subsequent rows for another string-CodePudding

I have the following DataFrame.

df = pd.DataFrame({'1': ['A','.','.','X','.','.'],
    '2':['.','.','.','.','A','.'],
    '3':['.','.','.','.','.','.'],
    '4':['.','.','.','.','.','X']})

I want to identify all instances where 'A' occurs and check to see if 'X' occurs within the next 3 rows.

After doing that I would like to execute a command based on these conditions.

an example of what I am trying to do would be...

for i, idx in df.iterrows():
    if idx == A:
        if X exists within next 3 rows:
            x= idx['1']
            y= idx['2']

Any help would be greatly appreciated.

CodePudding user response：

I am sure that the other answer could work if you were to explain what you really want to do. It would be more efficient as iterating over rows is slow.

However, here is a solution based on iterrows:

mask = df.eq('X').any(1)
mask = mask.where(mask).bfill(limit=3).fillna(False)

for idx, row in df.iterrows():
    if 'A' in row.values and mask[idx]:
        x = row['1']
        y = row['2']
        print(f'row {idx} matches: {x=}, {y=}')

example input (slightly different from yours):

   1  2  3  4
0  A  .  .  .
1  .  .  .  .
2  .  .  A  .
3  .  .  .  .
4  X  A  .  .
5  .  .  X  .

output:

row 2 matches: x='.', y='.'
row 4 matches: x='X', y='A'

CodePudding user response：

IIUC, you want to identify the cells where there is a value A and if within the next 3 rows, there is also a value X

I will use a more visual example for clarity (A/X/.):

   0  1  2  3  4  5
0  A  .  .  A  .  A
1  .  .  X  .  A  .
2  .  A  .  .  .  A
3  .  X  .  X  .  X
4  X  .  .  .  .  .

One can use eq to find the searched values and where bfill(limit=3) .fillna to extend the second mask to the previous lines.

# mask for the A
m1 = df.eq('A')

# mask for the X in the next 3 lines
m2 = df.eq('X')
m2 = m2.where(m2).bfill(limit=3).fillna(False)

# example of how to use the masks: replacing A with O
df[m1&m2] = 'O'

Example output:

   0  1  2  3  4  5
0  A  .  .  O  .  O
1  .  .  X  .  A  .
2  .  O  .  .  .  O
3  .  X  .  X  .  X
4  X  .  .  .  .  .

checking X for any column

Just change the second mask to:

m2 = df.eq('X').any(1)
m2 = m2.where(m2).bfill(limit=3).fillna(False)

output with this mask:

   0  1  2  3  4  5
0  O  .  .  O  .  O
1  .  .  X  .  O  .
2  .  O  .  .  .  O
3  .  X  .  X  .  X
4  X  .  .  .  .  .