基于多条件为DataFrame中几何对象Thing添加Side边标签的实现

我正在处理一个几何问题，其中名为Thing的对象（大多为矩形）会被切分为更小的矩形。我的DataFrame中包含s1和s2两列，数据示例如下：

Thing  s1  s2  id1  id2
0      A   2   4    0    1
1      A   2   4    1    2
2      A   2   4    1    3
3      A   2   4    0    4
4      B   6   2    1    1
5      B   6   2    2    0
6      B   6   2    3    0
7      B   6   2    4    1
8      B   6   2    5    1
9      B   6   2    6    0
10     C   6   2    1    4
11     C   3  10    2    0
12     C   3  10    3    6
13     C   3  10    4   67
14     C   3  10    5    0
15     C   3  10    5    3
16     C   3  10    6    4
17     C   3  10    6    4
18     C   3  10    6    3
19     C   3  10    6    3

规则说明 #

• 矩形判定：同一个Thing对应的所有行中，如果s1全为2且id1仅包含0/1两个取值，或s2全为2且id2仅包含0/1两个取值，则该Thing为矩形，否则为非矩形。
  • 示例：Thing A所有行s1=2且id1仅为0/1，属于矩形；Thing B所有行s2=2且id2仅为0/1，属于矩形。
• 边标识规则：
  • 对于s1全为2的矩形：id1=0对应边S1，id1=1对应边S2
  • 对于s2全为2的矩形：id2=0对应边S1，id2=1对应边S2
• 非矩形的所有行Side列统一标注为NaN

预期输出 #

需要新增Side列标注边，预期结果如下：

Thing  s1  s2  id1  id2 Side
0      A   2   4    0    1   S1
1      A   2   4    1    2   S2
2      A   2   4    1    3   S2
3      A   2   4    0    4   S1
4      B   6   2    1    1   S2
5      B   6   2    2    0   S1
6      B   6   2    3    0   S1
7      B   6   2    4    1   S2
8      B   6   2    5    1   S2
9      B   6   2    6    0   S1
10     C   6   2    1    4  NaN
11     C   3  10    2    0  NaN
12     C   3  10    3    6  NaN
13     C   3  10    4   67  NaN
14     C   3  10    5    0  NaN
15     C   3  10    5    3  NaN
16     C   3  10    6    4  NaN
17     C   3  10    6    4  NaN
18     C   3  10    6    3  NaN
19     C   3  10    6    3  NaN

现有问题 #

之前尝试了循环逐Thing处理和嵌套np.where两种方案，错误原因是直接按行判断s1/s2是否为2，没有先判定整个Thing是否为矩形，导致非矩形Thing中出现0/1的id1/id2时被错误标注。

核心思路是先按Thing分组完成矩形判定，再对已确认是矩形的分组进行边标注，避免非矩形数据被错误匹配。
使用pandas的分组转换和np.select实现如下：

import pandas as pd
import numpy as np

# 按Thing分组判定矩形类型
g = Named.groupby('Thing')
# s1类矩形：全组s1=2，且id1仅包含0/1
Named['is_s1_rect'] = g['s1'].transform(lambda x: (x == 2).all()) & g['id1'].transform(lambda x: x.isin([0,1]).all())
# s2类矩形：全组s2=2，且id2仅包含0/1
Named['is_s2_rect'] = g['s2'].transform(lambda x: (x == 2).all()) & g['id2'].transform(lambda x: x.isin([0,1]).all())

# 多条件匹配生成Side列
conditions = [
    Named['is_s1_rect'] & (Named['id1'] == 0),
    Named['is_s1_rect'] & (Named['id1'] == 1),
    Named['is_s2_rect'] & (Named['id2'] == 0),
    Named['is_s2_rect'] & (Named['id2'] == 1)
]
choices = ['S1', 'S2', 'S1', 'S2']
Named['Side'] = np.select(conditions, choices, default=np.nan)

# 删除中间辅助列
Named.drop(columns=['is_s1_rect', 'is_s2_rect'], inplace=True)

执行上述代码后即可得到与预期完全一致的结果。 #

内容的提问来源于stack exchange，提问作者Serge de Gosson de Varennes