如何通过中断执行嵌套列表理解？

Question

我有一个很大的距离数据框，我想对其进行分类。

df_norms = pd.DataFrame([[0, 100, 200, 4000000]], 
                        columns=['mode', 'min', 'medium', 'max'])

df_afst = pd.DataFrame([[0, 50, -1],
                        [0, 150, -1],
                        [0, 0, -1],
                        [0, 250, -1]], 
                       columns = ['train', 'station', 'bbh'])

规范 DataFrame 说对于项目零，每个距离 <= 100 被分类为 0，接下来当它 <= 200 时它是 1，最后是 catch-all <= 一个大数它是 2。

使用 for 循环很容易做到这一点。示例：

for i in [0]: # 1 element list just for the example
    bbh_id = i + 2
    mode = df_afst.iloc[0, i]
    for iy, y in enumerate(df_afst[df_afst.columns[i+1]].values):
        for ix, x in enumerate(df_norms.iloc[mode]):
            if x > y:
                df_afst.loc[iy, df_afst.columns[bbh_id]] = ix - 1
                break

之前：

   train  station  bbh
0      0       50   -1
1      0      150   -1
2      0        0   -1
3      0      250   -1

及之后

   train  station  bbh
0      0       50    0
1      0      150    1
2      0        0    0
3      0      250    2

我想在列表理解中执行此操作，但不知道如何执行此操作：break 使其难以执行。我能做的最好的是：

for i in [0]:
    bbh_id = i + 2
    mode = df_afst.iloc[0, i]
    r = [ix - 1 for iy, y in enumerate(df_afst[df_afst.columns[i+1]].values)
                     for ix, x in enumerate(df_norms.iloc[mode])
                         if x > y]

 # results in : [0, 1, 2, 1, 2, 0, 1, 2, 2]

如您所见，如果拆分结果，结果是正确的：

[0, 1, 2 | 1, 2 | 0, 1, 2 | 2]

我只需要子列表的第一个，不知道怎么做。我无法模拟 break。尝试了 min、[any][1] 和 next，但就是做错了。有人有什么想法吗？

更新

@chepner 正确地纠正了我的示例不一致。对不起。 @Thierry Lathuille 正确地指出列表理解并不总是正确的工具。他在这一点上是完全正确的，因为我不知道它们何时是正确的工具，这就是我想了解在这种情况下它是如何工作的原因。

我在这个答案上得到的两个答案对我很有启发。我从来没有听说过 pandas cut 也从来没有用过 numpy argwhere。

出于好奇，我做了一个小基准。

print('\n*** pd.cut')
cpu = time.time()
cuts = df_norms.iloc[0].tolist()
bbh3 = pd.cut(df_afst['station'], cuts, labels=False, include_lowest=True)
df_afst['bbh'] = bbh3
print('CPU {:.4f} seconds'.format(time.time() - cpu))
    
print('\n*** Using numpy and its functions')
cpu = time.time()
bbh2 = [np.min(np.argwhere(np.less(td, df_norms.values.ravel()))-1) for td in df_afst.station.values]
df_afst['bbh'] = bbh2
print('CPU {:.4f} seconds'.format(time.time() - cpu))

print('\n*** Simple loop')                
cpu = time.time()
for i in [0]:
    bbh_id = i + 2
    mode = df_afst.iloc[0, i]
    for iy, y in enumerate(df_afst[df_afst.columns[i+1]].values):
        for ix, x in enumerate(df_norms.iloc[mode]):
            if x > y:
                df_afst.loc[iy, df_afst.columns[bbh_id]] = ix - 1
                break

print('CPU {:.4f} seconds'.format(time.time() - cpu))

print('\n*** Wrong approach')                
cpu = time.time()
for i in [0]:
    bbh_id = i + 2
    mode = df_afst.iloc[0, i]
    r = [ix - 1 for iy, y in enumerate(df_afst[df_afst.columns[i+1]].values)
                     for ix, x in enumerate(df_norms.iloc[mode])
                         if x > y]
print('CPU {:.4f} seconds'.format(time.time() - cpu))

我将示例中的数据集从 4 扩大到 2,000,000，接近我的数据集 10,000,000。我得到的结果很有趣：

*** pd.cut
CPU 0.0131 seconds

*** Using numpy and its functions
CPU 29.4257 seconds

*** Simple loop
CPU 214.5378 seconds

*** Wrong approach
CPU 103.5768 seconds

pandas 剪切函数的加速效果令人难以置信。我仔细检查了结果，但似乎真的没问题。

两个答案，都正确而且很有见地。我决定将@carlos melus 的答案标记为正确答案，因为他最接近我要求的列表理解。

Answer 1

您可以使用 numpy 向量化计算：

[np.min(np.argwhere(np.less(td, df_norms.values.ravel()))-1) for td in df_afst.station.values]

np.less 将 df_afst.station 中的每个距离与 df_norms 和 returns 布尔矩阵中的所有值进行比较，如果 td 是小于df_norms.

中对应的值

例如，np.less(50, [0, 100, 200, 4000000]) returns: array([False, True, True, True])

使用 np.argwhere，我们提取输出数组中 True 值的索引，从 1 开始，因此我们减去 1 使其从 0 开始。从那里，获取数组中为 True 的最小索引，这就是您要查找的值。

您可以运行列表理解中的所有这些，结果将是：[0, 1, 0, 2]

Answer 2

使用 pd.cut() 你会受益匪浅:

假设您想要将 df_afst['station'] 中的值装箱（问题并不完全清楚，但我根据示例进行猜测），您可以这样做：

cuts = df_norms.iloc[0].tolist()
bbh = pd.cut(df_afst['station'], cuts, labels=False, include_lowest=True)

或更直接：

bbh = pd.cut(df_afst['station'], [-1, 100, 200, float('inf')], labels=False)

结果：

>>> bbh
0    0
1    1
2    0
3    2
Name: station, dtype: int64

当然，您也可以将其分配给列。

这将比 Python 循环（显式或理解）快几个数量级。

如何通过中断执行嵌套列表理解？

How to perform a nested list comprehension with a break?

python

list-comprehension

dataframe

pandas