python - 如何在pandas数据框中提取垂直条件下的一些行

Question

现在我有下面的数据框

    A B C   
1   a 0 6  
1   b 1 5
1   c 0 8
1   d 0 9
2   e 0 1
2   f 1 2 
2   g 0 3
3   h 0 4
3   i 0 5
3   j 0 6

首先，我想提取 df.B==1

   A B C
1  b 1 5
2  f 1 2

第二，我想提取 C 列编号周围的行 (+ -1，在本例中，行 c 0 8 紧邻 b 1 5，但 df.C=8 的编号，因此该列会被删除。如果 df.C=4or6 则会添加这一行)它们是否彼此相邻并不重要。

  A B C
1 a 0 6
1 b 1 5
2 e 0 1
2 f 1 2
2 g 0 3

我想获取上面的数据框。我怎样才能做到这一点？

首先我尝试

df[df.B==1]

但是，我不知道下一步......

Answer 1

# find where `B` is one
bs = np.where(df.B.values == 1)[0]

# union `bs` with + 1 and - 1
idx = np.unique(np.concatenate([bs - 1, bs, bs + 1]))

# don't go below zero or above last row
idx = idx[(0 <= idx) & (idx < len(df))]

# use `iloc` to index correct rows
df.iloc[idx]

   A  B  C
1  a  0  6
1  b  1  5
1  c  0  8
2  e  0  1
2  f  1  2
2  g  0  3