替换 numpy 数组中的值时防止字符串被截断

Question

假设我有数组a和b

a = np.array([1,2,3])
b = np.array(['red','red','red'])

如果我要对这些数组应用一些像这样的花哨索引

b[a<3]="blue"

我得到的输出是

array(['blu', 'blu', 'red'], dtype='<U3')

我知道这个问题是因为 numpy 最初只为 3 个字符分配空间，因此它无法将整个单词 blue 放入数组中，我可以使用什么解决方法？

目前我正在做

b = np.array([" "*100 for i in range(3)])
b[a>2] = "red"
b[a<3] = "blue"

但这只是一种解决方法，这是我的代码中的错误吗？或者是numpy的一些问题，我该如何解决这个问题？

Answer 1

您可以通过将dtypeof设置为b来处理可变长度字符串"object"：

import numpy as np
a = np.array([1,2,3])
b = np.array(['red','red','red'], dtype="object")

b[a<3] = "blue"

print(b)

这输出：

['blue' 'blue' 'red']

这dtype将处理字符串或其他通用 Python 对象。这也必然意味着在幕后您将拥有一个numpy指针数组，因此不要指望使用原始数据类型时获得的性能。