numpy 3d数组-沿给定轴选择最大元素

Question

我有一个尺寸为（x，y，z）=（5，50，4）的numpy数组。 对于每个（x，y）对，我想找到沿z轴的最大值的索引。 该索引在range（4）中。 我想选择所有这些“最大”元素并将它们设置为1。然后，我想选择所有其他元素并将它们设置为零。

为了以另一种方式解释它，我想查看“ z”方向上的所有向量（这些向量总共有x * y个）。 我想将最大元素设置为1，将所有其他元素设置为0。例如，向量（0.25，0.1，0.5，0.15）将变为（0，0，1，0）。

我尝试了很多不同的东西。 argmax函数似乎应该有所帮助。 但是如何使用它正确选择元素？ 我尝试过类似...

x = data
i = x.argmax(axis = 2)
x[i] # shape = (5, 50, 50, 4)
x[:,:,i] # shape = (5, 50, 5, 50)
x[np.unravel_index(i), x.shape] # shape = (5, 50)

最后一个使用np.unravel_index的形状正确，但是所选索引不是沿z轴的最大值。 所以我遇到了麻烦。 如果有人可以提供帮助，那将真的很棒。 谢谢！

---

编辑：这是我发现的一种方法。 但是，如果有人有更快的速度，请告诉我！

def fix_vector(a):
    i = a.argmax()
    a = a*0
    a[i] = 1
    return a

y = np.apply_along_axis(fix_vector, axis=2, arr=x)

我真的很想优化此方法，因为我多次调用了此函数。

---

编辑：感谢DSM为一个不错的解决方案。 这是注释中要求的一个小的示例数据集。

data = np.random.random((3,5,4))
desired_output = np.apply_along_axis(fix_vector, axis=2, arr=data)

这使用了我上面发布的fix_vector函数，但是DSM的解决方案更快。 再次感谢！

Answer 1

这不是特别优雅，但是：

def faster(x):
    d = x.reshape(-1, x.shape[-1])
    d2 = np.zeros_like(d)
    d2[np.arange(len(d2)), d.argmax(1)] = 1
    d2 = d2.reshape(x.shape)
    return d2

似乎比fix_vector方法快一点。 例如：

>>> x,y,z = 5,50,4
>>> data = np.random.random((x,y,z))
>>> np.allclose(orig(data), faster(data))
True
>>> %timeit -n 1000 orig(data)
1000 loops, best of 3: 5.77 ms per loop
>>> %timeit -n 1000 faster(data)
1000 loops, best of 3: 36.6 µs per loop

Answer 2

这可以通过numpy.where完成

import numpy as np
a = np.array([[[ 0.25,  0.10 ,  0.50 ,  0.15],
               [ 0.50,  0.60 ,  0.40 ,  0.30]],
              [[ 0.25,  0.50 ,  0.20 ,  0.70],
               [ 0.80,  0.10 ,  0.50 ,  0.15]]])

在最后一个轴上找到最大值

b = a.max(-1)    #shape is (2,2)

向b添加一个轴，以便它将在a 广播并创建一个布尔数组。

condition = a == b[..., np.newaxis]

使用numpy.where进行替换。

c = np.where(condition, 1, 0)

>>> c
array([[[0, 0, 1, 0],
        [0, 1, 0, 0]],

       [[0, 0, 0, 1],
        [1, 0, 0, 0]]])

---

def f(a):
    b = a.max(-1)
    condition = a == b[..., np.newaxis]
    return np.where(condition, 1, 0)

---

所以我不满意-我已经看到了reshape比添加新轴快的情况，所以我玩了一些。 似乎numpy.where本身有点慢。 在赋值中使用布尔索引会产生与@DSM几乎相同的性能。

def h(a):
    b = a.max(-1)
    condition = a == b[..., np.newaxis]
    a[condition] = 1
    a[np.logical_not(condition)] = 0
    return a

Answer 3

样本x ， (2,3,4)的形状：

In [934]: x
Out[934]: 
array([[[ 5,  7,  4, 15],
        [ 9, 23, 14,  9],
        [17, 13,  2,  1]],

       [[ 2, 18,  3, 18],
        [10, 12, 23, 14],
        [17, 23, 19, 13]]])

沿最后一个轴的max -足够清晰，一个（2,3）数组：

In [935]: x.max(2)
Out[935]: 
array([[15, 23, 17],
       [18, 23, 23]])

沿该轴的对应索引

In [936]: I=x.argmax(2)

In [937]: I
Out[937]: 
array([[3, 1, 0],
       [1, 2, 1]], dtype=int32)

用这些索引索引x一种方法。 此时，我正在手动构造0和1轴索引。 它需要自动化

In [938]: x[np.arange(2)[:,None], np.arange(3)[None,:], I]
Out[938]: 
array([[15, 23, 17],
       [18, 23, 23]])

相同的索引可用于设置另一个数组中的值：

In [939]: y=np.zeros_like(x)   
In [940]: y[np.arange(2)[:,None],np.arange(3)[None,:],I]=x.max(2)
In [941]: y
Out[941]: 
array([[[ 0,  0,  0, 15],
        [ 0, 23,  0,  0],
        [17,  0,  0,  0]],

       [[ 0, 18,  0,  0],
        [ 0,  0, 23,  0],
        [ 0, 23,  0,  0]]])

诀窍是想出一种更简化的方式来生成索引元组：

(np.arange(2)[:,None], np.arange(3)[None,:], I)

(array([[0],
    [1]]), 
 array([[0, 1, 2]]), 
 array([[3, 1, 0],
    [1, 2, 1]], dtype=int32))

np.ix_ ，尽管串联2个元组有点难看：

J = np.ix_(range(x.shape[0]), range(x.shape[1])) + (I,)

在扁平数组上建立索引可以更快：

In [998]: J1 = np.ravel_multi_index(J, x.shape)
In [999]: J1
Out[999]: 
array([[ 3,  5,  8],
     [13, 18, 21]], dtype=int32)

In [1000]: x.flat[J1]
Out[1000]: 
array([[15, 23, 17],
   [18, 23, 23]])

In [1001]: y.flat[J1]=2

In [1002]: y
Out[1002]: 
array([[[0, 0, 0, 2],
        [0, 2, 0, 0],
        [2, 0, 0, 0]],

       [[0, 2, 0, 0],
        [0, 0, 2, 0],
        [0, 2, 0, 0]]])

但是我仍然需要构造J索引元组。

这与DSM的解决方案相似，但坚持3维。 将尺寸的2展平可以消除xi_结构的需要。

---

def max_index(x, n, fn=np.argmax):
    # return multidimensional indexing tuple
    # applying fn along axis n
    I = fn(x, axis=n)
    alist = list(np.ix_(*[range(i) for i in I.shape]))
    alist.insert(n, I)
    return tuple(alist)

是用于生成类似J索引元组的通用函数

例如

y = np.zeros_like(x)
I = max_index(x, 1, np.argmin)
y[I] = x[I]
I = max_index(x, 1, np.argmax)
y[I] = x[I]

将x的最小值和最大值（沿轴1）插入y 。 速度大约是一半DSM小号faster 。