python - 使用 bool 数组在 DataFrame(或 ndarray)中有效设置一维范围值

Question

先决条件

import numpy as np
import pandas as pd

INPUT1: bool 二维数组(示例数组如下)

x = np.array(
    [[False,False,False,False,True],
     [True,False,False,False,False],
     [False,False,True,False,True],
     [False,True,True,False,False],
     [False,False,False,False,False]])

INPUT2:1D 范围值(示例如下)

y=np.array([1,2,3,4])

预期输出:2D ndarray

   [[0,0,0,0,1],
    [1,0,0,0,2],
    [2,0,1,0,1],
    [3,1,1,0,2],
    [4,2,2,0,3]]

我想有效地为 2d ndarray(INPUT1) 中的每个 True 设置一个范围值(垂直向量)。是否有一些有用的 API 或解决方案可用于此目的？

Answer 1

不幸的是我无法想出一个优雅的解决方案，所以我想出了多个不优雅的解决方案。我能想到的两种主要方法是

1. 对每个 True 值进行强力循环并分配切片，以及
2. 使用单个索引赋值来替换必要的值。

事实证明，这些方法的时间复杂度并不简单，因此根据数组的大小，任何一种方法都可以更快。

使用您的示例输入:

import numpy as np

x = np.array(
    [[False,False,False,False,True],
     [True,False,False,False,False],
     [False,False,True,False,True],
     [False,True,True,False,False],
     [False,False,False,False,False]])
y = np.array([1,2,3,4])
refout = np.array([[0,0,0,0,1],
    [1,0,0,0,2],
    [2,0,1,0,1],
    [3,1,1,0,2],
    [4,2,2,0,3]])

# alternative input with arbitrary size:
# N = 100; x = np.random.rand(N,N) < 0.2; y = np.arange(1,N)

def looping_clip(x, y):
    """Loop over Trues, use clipped slices"""
    nmax = x.shape[0]
    n = y.size

    # initialize output
    out = np.zeros_like(x, dtype=y.dtype)
    # loop over True values
    for i,j in zip(*x.nonzero()):
        # truncate right-hand side where necessary
        out[i:i+n, j] = y[:nmax-i]
    return out

def looping_expand(x, y):
    """Loop over Trues, use an expanded buffer"""
    n = y.size
    nmax,mmax = x.shape
    ivals,jvals = x.nonzero()

    # initialize buffed-up output
    out = np.zeros((nmax + max(n + ivals.max() - nmax,0), mmax), dtype=y.dtype)
    # loop over True values
    for i,j in zip(ivals, jvals):
        # slice will always be complete, i.e. of length y.size
        out[i:i+n, j] = y
    return out[:nmax, :].copy() # rather not return a view to an auxiliary array

def index_2d(x, y):
    """Assign directly with 2d indices, use an expanded buffer"""
    n = y.size
    nmax,mmax = x.shape
    ivals,jvals = x.nonzero()

    # initialize buffed-up output
    out = np.zeros((nmax + max(n + ivals.max() - nmax,0), mmax), dtype=y.dtype)

    # now we can safely index for each "(ivals:ivals+n, jvals)" so to speak
    upped_ivals = ivals[:,None] + np.arange(n) # shape (ntrues, n)
    upped_jvals = jvals.repeat(y.size).reshape(-1, n) # shape (ntrues, n)

    out[upped_ivals, upped_jvals] = y # right-hand size of shape (n,) broadcasts

    return out[:nmax, :].copy() # rather not return a view to an auxiliary array

def index_1d(x,y):
    """Assign using linear indices, use an expanded buffer"""
    n = y.size
    nmax,mmax = x.shape
    ivals,jvals = x.nonzero()

    # initialize buffed-up output
    out = np.zeros((nmax + max(n + ivals.max() - nmax,0), mmax), dtype=y.dtype)

    # grab linear indices corresponding to Trues in a buffed-up array
    inds = np.ravel_multi_index((ivals, jvals), out.shape)

    # now all we need to do is start stepping along rows for each item and assign y
    upped_inds = inds[:,None] + mmax*np.arange(n) # shape (ntrues, n)

    out.flat[upped_inds] = y  # y of shape (n,) broadcasts to (ntrues, n)

    return out[:nmax, :].copy() # rather not return a view to an auxiliary array


# check that the results are correct
print(all([np.array_equal(refout, looping_clip(x,y)),
           np.array_equal(refout, looping_expand(x,y)),
           np.array_equal(refout, index_2d(x,y)),
           np.array_equal(refout, index_1d(x,y))]))

我尝试记录每个函数，但这里有一个概要:

1. looping_clip 循环输入中的每个 True 值，并分配给输出中相应的切片。当切片的一部分沿第一维超出数组边缘时，我们会在右侧注意缩短分配的数组。
2. looping_expand 循环输入中的每个 True 值，并在分配填充的输出数组后分配给输出中相应的 full 切片，确保让每一片都充满。当分配更大的输出数组时，我们会做更多的工作，但我们不必在赋值时缩短右侧。我们可以在最后一步中省略 .copy() 调用，但我不喜欢返回一个非平凡的跨步数组(即辅助数组的 View 而不是正确的副本)，因为这可能会导致给用户带来不为人知的惊喜。
3. index_2d 计算要分配的每个值的 2d 索引，并假设重复索引将按顺序处理。这是无法保证的! (稍后会详细介绍。)
4. index_1d 使用线性化索引并索引到输出的 flatiter 中，执行相同的操作。

以下是使用随机数组的上述方法的计时(请参阅开头附近的注释行):

我们可以看到，对于小型和大型数组，循环版本更快，但对于大约 10 到 150 之间的线性大小，索引版本更好。我没有采用更高尺寸的原因是索引案例开始使用大量内存，而且我不想担心这种时间困惑。

为了使上述情况变得更糟，请注意，索引版本假定按顺序处理奇特索引场景中的重复索引，因此当处理 True 值时，这些值在数组中“较低” ，以前的值将根据您的要求被覆盖。只有一个问题:this is not guaranteed :

For advanced assignments, there is in general no guarantee for the iteration order. This means that if an element is set more than once, it is not possible to predict the final result.

这听起来不太令人鼓舞。虽然在我的实验中，索引似乎是按顺序处理的(根据 C 顺序)，但这也可能是巧合，或者是实现细节。因此，如果您想使用索引版本，请确保在您的特定版本以及特定尺寸和形状上这仍然适用。

我们可以通过自己删除重复索引来使分配更安全。为此，我们可以利用this answer by Divakar关于相应的问题:

def index_1d_safe(x,y):
    """Same as index_1d but use Divakar's safe solution for reducing duplicates"""
    n = y.size
    nmax,mmax = x.shape
    ivals,jvals = x.nonzero()

    # initialize buffed-up output
    out = np.zeros((nmax + max(n + ivals.max() - nmax,0), mmax), dtype=y.dtype)

    # grab linear indices corresponding to Trues in a buffed-up array
    inds = np.ravel_multi_index((ivals, jvals), out.shape)

    # now all we need to do is start stepping along rows for each item and assign y
    upped_inds = inds[:,None] + mmax*np.arange(n) # shape (ntrues, n)

    # now comes https://stackoverflow.com/a/44672126
    # need additional step: flatten upped_inds and corresponding y values for selection
    upped_flat_inds = upped_inds.ravel() # shape (ntrues, n) -> (ntrues*n,)
    y_vals = np.broadcast_to(y, upped_inds.shape).ravel() # shape (ntrues, n) -> (ntrues*n,)

    sidx = upped_flat_inds.argsort(kind='mergesort')
    sindex = upped_flat_inds[sidx]
    idx = sidx[np.r_[np.flatnonzero(sindex[1:] != sindex[:-1]), upped_flat_inds.size-1]]
    out.flat[upped_flat_inds[idx]] = y_vals[idx]

    return out[:nmax, :].copy() # rather not return a view to an auxiliary array

这仍然会重现您的预期输出。问题是现在该函数需要更长的时间才能完成:

真糟糕。考虑到我的索引版本仅对于中间数组大小而言更快，以及它们的更快版本如何不能保证工作，也许最简单的方法是仅使用其中一个循环版本。当然，这并不是说我没有错过任何最佳矢量化解决方案。