基于Pandas DataFrame列约束规则生成新列col3的技术问询

针对你给出的DataFrame和规则，这里有两种实用的方法来实现col3的生成，都能精准满足需求：

方法1：循环分组处理（直观易懂） #

这种方法逻辑清晰，适合刚接触pandas的朋友理解：

import pandas as pd

# 构建你的输入DataFrame
df = pd.DataFrame({
    'ID': [1, 1, 1, 1, 1, 1],
    'col1': ['A', 'A', 'B', 'B', 'C', 'D'],
    'col2': [1, 2, 3, 4, 5, 6]
})

# 先给col3设置默认值1
df['col3'] = 1

# 按ID和col1分组处理每个组
for _, group in df.groupby(['ID', 'col1']):
    # 只有当组内有多个col2实例时才需要调整
    if len(group) > 1:
        # 先按col2排序（确保顺序正确，避免原始数据无序的情况）
        sorted_group = group.sort_values('col2')
        # 计算当前col2与前一个col2的差值
        col2_diff = sorted_group['col2'].diff()
        # 找到差值为1的行，将对应位置的col3设为2
        target_indices = sorted_group[col2_diff == 1].index
        df.loc[target_indices, 'col3'] = 2

# 查看结果
print(df)

执行后输出的结果如下：

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

逻辑解释： #

• 先给所有行的col3赋值默认值1，满足规则的基础要求；
• 按ID和col1分组，确保我们只在同一ID、同一col1的范围内处理数据；
• 对于每组，只有当组内存在多个col2值时才进行调整；
• 通过diff()方法计算相邻col2的差值，当差值为1时，说明当前行的col2是前一行+1，此时将col3设为2。

方法2：矢量化操作（高效简洁） #

如果你的数据集较大，推荐用这种矢量化方法，避免循环，性能更优：

import pandas as pd

df = pd.DataFrame({
    'ID': [1, 1, 1, 1, 1, 1],
    'col1': ['A', 'A', 'B', 'B', 'C', 'D'],
    'col2': [1, 2, 3, 4, 5, 6]
})

# 设置默认值1
df['col3'] = 1

# 分组计算每个组内col2的前后差值
df['col2_diff'] = df.groupby(['ID', 'col1'])['col2'].diff()

# 找到差值为1的行，将col3设为2
df.loc[df['col2_diff'] == 1, 'col3'] = 2

# 可以删除临时的差值列（如果不需要的话）
df = df.drop('col2_diff', axis=1)

print(df)

这个方法的结果和方法1完全一致，但代码更简洁，处理大数据时速度会快很多，因为pandas的矢量化操作是底层优化过的。

逻辑解释： #

• 同样先设置col3默认值1；
• 利用groupby+diff()直接生成每个组内col2的相邻差值，这一步是矢量化操作，无需手动循环；
• 直接通过布尔索引定位到差值为1的行，修改对应的col3值即可。

内容的提问来源于stack exchange，提问作者babz