使用元组进行多索引-解网

问：

我有一个多维的.我知道第一个 N 维和最后一个 M 维的形状。例如，np.array

>>> n = (3,4,5)
>>> m = (6,)
>>> a = np.ones(n + m)
>>> a.shape
(3, 4, 5, 6)

使用元组作为索引可以快速索引，例如前 N 个维度，例如

>>> i = (1,1,2)
>>> a[i].shape
(6,)

使用 list 不会给我我需要的相同结果

>>> i = [1,1,2]
>>> a[i].shape
(3, 4, 5, 6)

但是我在进行多索引（检索/分配值）时遇到了麻烦。例如

>>> i = (1,1,2)
>>> j = (2,2,2)

我需要传递类似的东西

>>> a[[i, j]]

并得到的输出形状。(2, 6)

相反，我得到了

>>> a[[i, j]].shape
(2, 3, 4, 5, 6)

或

>>> a[(i, j)].shape
(3, 5, 6)

我总是可以循环或更改索引事物的方式（例如使用和），但是有没有更 pythonic 的方法来实现我需要的东西？np.reshapenp.unravel_index

编辑对于任意数量的索引，我都需要它，例如，

>>> i = (1,1,2)
>>> j = (2,2,2)
>>> k = (0,0,0)
...

python numpy 多索引

import time
import numpy as np

n = (3, 4, 5)
m = (2, 4, 6)
a = np.random.rand(*(n + m))

indices = [tuple(np.random.randint(n)) for _ in range(100)]

test = 100000

t1 = time.time()

for _ in range(test):
  x = a[tuple(zip(*indices))]  

t2 = time.time()

for _ in range(test):
  y = np.array([a[i] for i in indices])

t3 = time.time()

for _ in range(test):
  b = a.reshape(np.prod(n), np.prod(m))
  j = [np.ravel_multi_index(i, n) for i in indices]
  z = b[j].reshape(-1, *m)

t4 = time.time()

>>> print(t2 - t1, t3 - t2, t4 - t3)
>>> print(x.shape, y.shape, z.shape)
>>> assert(np.all(x - y == 0))
>>> assert(np.all(x - z == 0))
3.8938186168670654 4.831823110580444 37.469303131103516
(100, 2, 4, 6) (100, 2, 4, 6) (100, 2, 4, 6)

1赞 Chrysophylaxs 11/11/2023 #4

考虑以下指数列表：

idx = [
    (1, 1, 2),  # Your i
    (2, 2, 2),  # Your j
    (0, 0, 0),  # Your k
    (1, 2, 1),  # ... 
    (2, 0, 1),  # extend as necessary
]

和你的数组形状 .a(3, 4, 5, 6)

当你写的时候，numpy是这样解释的：out = a[idx]

out = np.array([
    [a[1], a[1], a[2]],
    [a[2], a[2], a[2]],
    [a[0], a[0], a[0]],
    [a[1], a[2], a[1]],
    [a[2], a[0], a[1]],
])

其中，例如，只是的第一个子数组，因此具有形状！a[0]a(4, 5, 6)

因此，您只剩下表示形状数组（索引的形状）的内容，其中包含！(...最终结果为，或）。(5, 3)(4, 5, 6)a(5, 3, 4, 5, 6)np.shape(idx) + a.shape[1:]

相反，您想要的是以下内容：

out = np.array([
    a[1, 1, 2],
    a[2, 2, 2],
    a[0, 0, 0],
    a[1, 2, 1],
    a[2, 0, 1],
])

在 numpy 中“矢量化”完成它的方法如下：

out = a[
    [1, 2, 0, 1, 2],  # [idx[0][0], idx[1][0], idx[2][0], ...]
    [1, 2, 0, 2, 0],  # [idx[0][1], idx[1][1], idx[2][1], ...]
    [2, 2, 0, 1, 1]   # [idx[0][2], idx[1][2], idx[2][2], ...]
]

该行为记录在索引指南中：

高级索引始终作为一个索引进行广播和迭代：
result[i_1, ..., i_M] == x[ind_1[i_1, ..., i_M], ind_2[i_1, ..., i_M],
                           ..., ind_N[i_1, ..., i_M]]

若要将原始文件转换为此类索引器，可以使用此技巧。idxtuple(zip(*idx))

Numpy的索引系统具有神奇的灵活性，但这种灵活性的代价是这些“简单”的任务变得不直观......至少在我看来;)

上一个：将 .loc 与 DataFrame 的级别子集一起使用

下一个：Python Pandas - 重复字典以适应 MultiIndex Dataframe

使用元组进行多索引

Multi-indexing with tuples

评论

评论

评论