如何从第三个数组给出的对应关系将一个 numpy 标签数组转换为一个颜色数组？

Question

我需要帮助找到一种有效的方法（尽可能快）将 numpy 标签数组转换为 numpy 颜色数组。

举个简单的例子：

A 包含标签（整数）：
```
A = [0,45,45,22,0,45,45,22]
```
B 包含所有标签：
```
B = np.unique(A) = [0,45,22]
```
C 包含 RGB 值：
```
C = [[1,0,0], [0,1,0], [0,0,1]]
```

而C的第i个元素是B中第i个标签的颜色。例如，标签 45 的颜色是 [0,1,0].

据此，A 应转换为：

[[1,0,0], [0,1,0], [0,1,0], [0,0,1], ...]

我已经尝试过下面的代码，但是速度很慢：

result = np.array([C[np.where(B==x)[0][0]] for x in A])

有人知道更有效的解决方案吗？

提前致谢:)

Answer 1

您可以使用 np.unique 的反向索引：

import numpy as np

A = np.array([0,45,45,22,0,45,45,22])
C = np.array([[1,0,0], [0,1,0], [0,0,1]])
_, inverse_idx = np.unique(A, return_inverse=True)

result = C[inverse_idx]
# array([[1, 0, 0], [0, 0, 1], [0, 0, 1], [0, 1, 0], [1, 0, 0], [0, 0, 1], [0, 0, 1], [0, 1, 0]])

重要说明：np.unique returns 值和索引作为排序数组，因此 np.unique(A) 给出 [0, 22, 35]，而不是 [0, 45, 22]。如果您确实希望按照它们出现的顺序使用它，则需要使用 A:

值的原始索引进行额外的操作

import numpy as np

A = np.array([0,45,45,22,0,45,45,22])
C = np.array([[1,0,0], [0,1,0], [0,0,1]])
_, idx, inverse_idx = np.unique(A, return_index=True, return_inverse=True)

result = C[idx.argsort()[inverse_idx]]
# array([[1, 0, 0], [0, 1, 0], [0, 1, 0], [0, 0, 1], [1, 0, 0], [0, 1, 0], [0, 1, 0], [0, 0, 1]])

Answer 2

A = np.array([0,45,45,22,0,45,45,22])
B = np.unique(A)
C = np.array([[1,0,0], [0,1,0], [0,0,1]])


from numba import njit

@njit
def f(arr, labels, colors):
    result = np.zeros((len(arr), 3))
    for i, label in enumerate(labels):
        result[arr==label] = colors[i]
    return result

使用 A:

中的单个元素编译函数

f(A[:1], B, C)

现在：

result = f(A, B, C)

我的机器需要 9.5367431640625e-05 秒，而你的解决方案需要 3.123283e-04 秒

我还在 1'000'000 个数字上尝试了我的函数 A 并且它需要 0.0359194278717041 vs 5.217064619064331 秒你的解决方案

如何从第三个数组给出的对应关系将一个 numpy 标签数组转换为一个颜色数组？

How can I transform a numpy array of labels to an array of colors, from correspondance given by third array?

python

arrays

numpy