Python/Dataframe逻辑在具有限制的列上运行

当前df

ID col1 col2
1  A  1
1  A  2
1  B  3
1  B  4
1  C  5
1  D  6

我需要创建col3，它将是1或2，具体取决于以下条件： col3=1，除非对于同一ID和同一col1，存在多个col2实例。如果有多个实例，从col2的最低实例开始，分配col3=1，然后转到col2的下一个实例，如果下一个实例是前一个实例+1，则col3=2，否则col3=1

听起来令人困惑，但期望的结果是

ID col1 col2 col3
1  A    1    1
1  A    2    2
1  B    3    1
1  B    4    2
1  C    5    1
1  D    6    1

编辑

添加了ID=1，col1=A，col2=3（因为col2等于1+它上面的行），col3=2

---

添加了ID=1，col1=B，col2=17（因为col2不等于1+它上面的行），col3=1让我们使用

cumcount

：

df.assign(col3 = df.groupby(['ID','col1']).cumcount() + 1)

输出：

   ID col1  col2  col3
0   1    A     1     1
1   1    A     2     2
2   1    B     3     1
3   1    B     4     2
4   1    C     5     1
5   1    D     6     1

   ID col1  col2  col3  col3a
0   1    A     1     1      1
1   1    A     2     2      2
2   1    A     3     2      2
3   1    B     3     1      1
4   1    B     4     2      2
5   1    B    17     1      1
6   1    C     5     1      1
7   1    D     6     1      1

编辑以处理“剪裁”使用df.clip进行比较创建新列“col3a”： 输出：

   ID col1  col2  col3
0   1    A     1     1
1   1    A     2     2
2   1    B     3     1
3   1    B     4     2
4   1    C     5     1
5   1    D     6     1

   ID col1  col2  col3  col3a
0   1    A     1     1      1
1   1    A     2     2      2
2   1    A     3     2      2
3   1    B     3     1      1
4   1    B     4     2      2
5   1    B    17     1      1
6   1    C     5     1      1
7   1    D     6     1      1

或许

df.groupby('col1').col2.apply(lambda x : (x.diff().eq(1))).map({True:2,False:1})
Out[1148]: 
0    1
1    2
2    2
3    1
4    2
5    1
6    1
7    1
Name: col2, dtype: int64

---

请展示您迄今为止所做的尝试。谢谢@Scott。如果有另一行的ID1 col1=A col2=3，我想这个方法会失败。这里col3应该等于2而不是3？@babz您能用问题中的新行以及预期输出创建新数据集吗。我对你的问题有点困惑。向OP@babz这是一个更好的解决方案。@Wen，你熟悉pyspark吗？我正在寻找一个类似的解决方案，但在pyspark中没有熊猫。