时间戳

首页 > 解决方案 > 如何通过对索引向量的广播进行切片来修改 np 数组？

python - 如何通过对索引向量的广播进行切片来修改 np 数组？

问题描述

我有这个 mxn numpy 数组，我想对行元素应用某些操作。虽然，它必须只对那些索引在索引向量上的条目指定的那些元素之前进行强制转换。

我已经经历了经典的 for 循环方式，但我期待更多的 NumPythonic。

以下代码将完成这项工作：

for i,j in enumerate(x):
    M[i, 0:j] = 2*M[i, 0:j]

但我一直在寻找一种广播，禁止的方法。有任何想法吗？

例如，让我们说

M = [[ 1, 2, 3, 4, 5],
     [ 6, 7, 8, 9,10]]

x = [2, 3]

我们的应用是将某些元素加倍。根据 x 中指定的索引，我们应该得到结果数组：

M = [[ 2, 4, 3, 4, 5],
     [12,14,16, 9,10]]

标签： pythonnumpyarray-broadcasting

解决方案

以下是两种相关的方式，它们的速度大致相同：

import numpy as np
from timeit import timeit

M = [[ 1, 2, 3, 4, 5],
     [ 6, 7, 8, 9,10]]

x = [2, 3]

def f():
    MM = np.array(M)
    xx = np.array(x)
    MM[np.arange(5)<xx[:,None]] *= 2
    return MM

def g():
    MM = np.array(M)
    xx = np.array(x)
    MM *= 1 + (np.arange(5)<xx[:,None])
    return MM

print(f())
print(g())
M = 1000*M
x = 1000*x
print(timeit(f,number=1000))
print(timeit(g,number=1000))

样品运行：

[[ 2  4  3  4  5]
 [12 14 16  9 10]]
[[ 2  4  3  4  5]
 [12 14 16  9 10]]
1.1994759310036898
1.1547658089984907

推荐阅读