Numpy：用越来越多的数字填充条件子数组答案

【问题标题】：Numpy: fill conditional subarray with increasing numbersNumpy：用越来越多的数字填充条件子数组
【发布时间】：2021-12-18 13:34:35
【问题描述】：

我经常遇到这样的成语：say I have data like

N = 20                             # or some other number
a = np.random.randint(0, 10, N)    # or any other 1D np.array
predicate = lambda x: x%2 == 0     # or any other predicate

我遇到的成语大同小异

b = np.full_like(a, -1)
i1 = 0
for i, x in enumerate(a):
    if predicate(x):
        b[i] = i1
        i1 += 1

如何将其翻译为 numpy？以下：

b = np.full_like(a, -1)
m = some_predicate(a)
b[m] = np.arange(np.count_nonzero(m))

在我看来有点奇怪：对于这样一个简单的任务，这只有三行。特别是，我需要存储m，这让我感到不安，因为我需要引用它两次（因为我无法说“根据需要设置尽可能多的值”）。

【问题讨论】：

"arange with as many values as required" 是一个复杂的事情，需要某种生成器，这不像你在numpy 标准函数中找到的那样.
没有什么能阻止你在两行中这样做，b = ...; b[some_predicate(a)] = np.arange(some_predicate(a).shape[1])，但这真的很可怕。我不相信“简单任务”
顺便说一下，a 的定义是一个数字，而不是一维数组。即使对于一维数组， m 也将具有相同的形状，因此 m.shape[1] 没有意义。但是由于a 是一维数组，您的实际意图似乎应该是b[m] = np.arange(np.count_nonzero(m))。
@Reti43 建议的行（你已经称之为可怕的）在a 上对谓词进行两次评估。我想我不明白您在第一条评论的其余部分中所说的内容；你能换一种说法吗？关于你的第二个帖子：你是对的；改变了。
我觉得你对一个简单的任务应该有多少行感到困惑。顺便说一句，这不是一个简单的任务，它可以在 3 行中完成，这很优雅。为什么在存储对需要多次使用的东西的引用时会感到不安？第一种方法：python 级别的循环很慢，你调用一个函数 N 次。第二种方法：一次对整个数组应用一个调用函数和 numpy 操作 = 快速。

标签： python numpy numpy-ndarray

【解决方案1】：

Walrus 运算符的救援（从 Python 3.8 开始）：

i = -1
b = np.array([ -1 if not predicate(val) else (i := i+1) for val in a ])

或者（对于大型数组可能要快得多）

b = np.full_like(a, -1)
b[sel] = np.arange(np.count_nonzero(sel := predicate(a)))

【讨论】：