Python: 根据掩码将列插入Numpy数组

Question

假设我有以下数据：

mask = [[0, 1, 1, 0, 1]] # 2D mask
ip_array = [[4, 5, 2],
            [3, 2, 1],
            [1, 8, 6]] # 2D array

我想在ip_array中的每个掩码为0的位置插入0列。因此输出应该是这样的：

[[0, 4, 5, 0, 2]
 [0, 3, 2, 0, 1]
 [0, 1, 8, 0, 6]]

我是numpy函数的新手，正在寻找一种高效的方法来完成这个任务。感谢您的帮助！

- Animeartist

---

高效如何？不适用于Python循环吗？（只能使用NumPy吗？） - Vincent Bénet

我正在尝试使用numpy来完成这个。 - Animeartist

你的n值是多少？（mask和ip_array的长度平均值是多少？）最小值和最大值是多少？因为在“设置”阶段使用numpy会耗费时间，所以有一个阈值，在这个阈值之下使用它是一个好主意还是不好的主意... - Vincent Bénet

掩码和ndarray的长度会不时变化。 - Animeartist

Answer 1

以下是两个步骤实现的方法：

(i) 创建一个形状正确的零数组（第一个维度来自 ip_array，第二个维度来自 mask）

(ii) 在第二个维度上使用 mask 作为布尔掩码，并将 ip_array 的值分配给零数组。

out = np.zeros((ip_array.shape[0], mask.shape[1])).astype(int)
out[..., mask[0].astype(bool)] = ip_array
print(out)

输出：

[[0 4 5 0 2]
 [0 3 2 0 1]
 [0 1 8 0 6]]

- user7864386

Answer 2

这里有另一种方法，使用切片和 cumsum 掩码以及输入中的额外 0 列。cumsum 掩码将具有 ip_array +1 的索引，以及每当添加零时的 0。连接的数组具有额外的初始零列，因此使用 0 进行索引会产生一列零。

m = (mask.cumsum()*mask)[0]
# array([0, 1, 2, 0, 3])

np.c_[np.zeros(ip_array.shape[0]), ip_array][:,m].astype(int)

# array([[0, 4, 5, 0, 2],
#        [0, 3, 2, 0, 1],
#        [0, 1, 8, 0, 6]])

- mozway

---

我认为我找到了一种更简单的方法，但只需要三行代码。 - Mad Physicist

Answer 3

一个带参数的解决方案，比绿色勾选的方式更易理解。只有最后一行对操作很重要。

import numpy
import random

n1 = 5
n2 = 5
r = 0.7
random.seed(1)
a = numpy.array([[0 if random.random() > r else 1 for _ in range(n1)]])
n3 = numpy.count_nonzero(a)
b = numpy.array([[random.randint(1,9) for _ in range(n3)] for _ in range(n2)])
c = numpy.zeros((n2, n1))
c[:, numpy.where(a)[1]] = b[:]

结果：

a = array([[1, 0, 0, 1, 1]])
b = array([[8, 8, 7],
       [4, 2, 8],
       [1, 7, 7],
       [1, 8, 5],
       [4, 2, 6]])
c = array([[8., 0., 0., 8., 7.],
       [4., 0., 0., 2., 8.],
       [1., 0., 0., 7., 7.],
       [1., 0., 0., 8., 5.],
       [4., 0., 0., 2., 6.]])

在这里，您的时间处理取决于n值：

使用以下代码：

import numpy
import random
import time
import matplotlib.pyplot as plt

n1 = 5
n2 = 5
r = 0.7


def main(n1, n2):
    print()
    print(f"{n1 = }")
    print(f"{n2 = }")
    random.seed(1)
    a = numpy.array([[0 if random.random() > r else 1 for _ in range(n1)]])
    n3 = numpy.count_nonzero(a)
    b = numpy.array([[random.randint(1,9) for _ in range(n3)] for _ in range(n2)])
    t0 = time.time()
    c = numpy.zeros((n2, n1))
    c[:, numpy.where(a)[1]] = b[:]
    t = time.time() - t0
    print(f"{t = }")
    return t


t1 = [main(10**i, 10) for i in range(1, 8)]
t2 = [main(10, 10**i) for i in range(1, 8)]

plt.plot(t1, label="n1 time process evolution")
plt.plot(t2, label="n2 time process evolution")

plt.xlabel("n-values (log)")
plt.ylabel("Time processing (s)")
plt.title("Insert columns into a numpy array based on mask")
plt.legend()
plt.show()

- Vincent Bénet

Answer 4

mask = np.array([0, 1, 1, 0, 1])
#extract indices of zeros
mask_pos = (list(np.where(mask == 0)[0]))
ip_array =np.array([[4, 5, 2],
        [3, 2, 1],
        [1, 8, 6]])

#insert 0 at respextive mask position
for i in mask_pos:
    ip_array = np.insert(ip_array,i,0,axis=1)

print(ip_array)

- Mahantesh

Answer 5

可以说最简单的解决方案就是使用np.insert为您创建新列：

idx = np.flatnonzero(~np.array(mask[0], bool))
idx -= np.arange(len(idx))
np.insert(ip_array, idx, 0, axis=1)

从idx中减去np.arange(len(idx))是必要的，因为您要插入的数组还没有新列，所以旧数组中的索引会减少前面插入的列数。

简而言之：

idx = np.flatnonzero(~np.array(mask[0], bool))
np.insert(ip_array, idx - np.arange(len(idx)), 0, axis=1)

一行代码使用海象运算符（Python 3.8+）：

np.insert(ip_array, (idx := np.flatnonzero(~np.array(mask[0], bool))) - np.arange(len(idx)), 0, axis=1)

还有一个版本，但具有更多的冗余：

np.insert(ip_array, np.flatnonzero(~np.array(mask[0], bool)) - np.arange(len(mask[0]) - np.count_nonzero(mask[0])), 0, axis=1)

- Mad Physicist