时间戳

首页 > 解决方案 > 如何为熊猫数据框中的一行中的连续零构建掩码（真或假）？

python - 如何为熊猫数据框中的一行中的连续零构建掩码（真或假）？

问题描述

我有一个数据框并想创建一个掩码，标记其中出现连续零（3 个或更多）。

我想将零的出现标记为 True 当它们连续发生三次或更多次时。也就是说，连续三个或更多零，没有任何类型的中断（例如其他数字或 NaN）。因此，数组 [0,0,0,1,2,3,4] 变为 [True, True, True, False, False, False, False] 并且数组 [0,1,0,2,0,0 ,0] 变为 [False,False,False,False,True,True,True]。

下面的示例更好地说明了输入和输出。

输入：

col1    col2    col3    col4    col5
0       0       0       0       0
0       0       5       0       5
0       0       0       3       3
0       0       0       NaN     0

我希望结果类似于

col1  col2  col3  col4  col5
True  True  True  True  True
False False False False False
True  True  True  False False
True  True  True  False False

目前我正在对行进行迭代，但速度有点慢（这个数据框实际上有超过 100 万行）。

有什么方法可以避免在行上出现 for 循环（iterrows、apply 等）吗？

谢谢！

标签： pythonpandasnumpy

解决方案

我将展示一个 numpy 解决方案，并留给您转换为 pandas。

a = np.array([
    [0, 0, 0],
    [0, 0, 0],
    [0, 5, 0],
    [0, 0, 3],
    [0, 5, 3]]).T

首先，掩码数组并在每一行上用零填充掩码：

z = np.zeros((a.shape[0], 1), dtype=bool)
mask = np.concatenate((z, a == 0, z), axis=1)

通过查找掩码更改值的位置，您可以找到每次运行的长度。上一步中的填充确保第一个更改为“on”，最后一个更改为“off”：

locs = np.nonzero(np.diff(mask, axis=1))

现在是稍微棘手的部分。通常你会在 1D 中执行此操作，但这里是 2D。但是，填充可确保您的运行长度仍然准确：

run_lengths = locs[1][1::2] - locs[1][::2]

现在你想关闭mask那些代表少于元素的运行的n=3元素。有了您拥有的信息，直接创建新蒙版可能会更容易。

首先屏蔽太短的运行：

valid_runs = np.flatnonzero(run_lengths >= 3)

然后将输出设为 a int8，我们稍后将其视为 a bool：

result = np.zeros(a.shape, dtype=np.int8)

将每个有效运行的第一个元素设置为 1，将末尾的一个元素设置为 -1：

v = 2 * valid_runs
result[locs[0][v], locs[1][v]] = 1
v += 1
v = v[locs[1][v] < result.shape[1]]
result[locs[0][v], locs[1][v]] = -1

您可以就地获取结果的累积总和，并将其视为布尔掩码（因为int8和bool_具有相同的大小）：

result = np.cumsum(result, axis=1, out=result).view(bool)

最终结果是：

array([[ True,  True,  True,  True,  True],
       [False, False, False, False, False],
       [ True,  True,  True, False, False]])

TL;博士

这是一个完全通用的解决方案，可以在多维数组中的任何轴上工作，可以任意选择最小运行长度 ( n) 和数量 ( k)：

def mask_consecutive(a, k=0, n=3, axis=-1):
    shape = list(a.shape)
    shape[axis] = 1;
    z = np.zeros(shape, dtype=bool)
    mask = np.concatenate((z, a == k, z), axis=axis)
    locs = np.argwhere(np.diff(mask, axis=axis))
    run_lengths = locs[1::2, axis] - locs[::2, axis]
    valid_runs = np.flatnonzero(run_lengths >= n)
    result = np.zeros(a.shape, dtype=np.int8)
    v = 2 * valid_runs
    result[tuple(locs[v, :].T)] = 1
    v += 1
    v = v[locs[v, axis] < result.shape[axis]]
    result[tuple(locs[v, :].T)] = -1
    return np.cumsum(result, axis=axis, out=result).view(bool)

推荐阅读