根据字典替换尴尬数组中的值的有效方法?

Question

我有一个包含整数键和浮点值的字典。 我还有一个带有整数条目的 2D 笨拙数组(我使用的是 awkward1)。 我想根据字典将这些整数替换为相应的浮点数，保持尴尬的数组格式。

假设键从 0 到 999 运行，到目前为止我的解决方案是这样的：

resultArray = ak.where(myArray == 0, myDict.get(0), 0)
for key in range(1,1000):
resultArray = resultArray + ak.where(myArray == key, myDict.get(key), 0)

有没有更快的方法来做到这一点？

更新

我的工作代码的最小可重现示例：

import awkward as ak # Awkward 1
myArray = ak.from_iter([[0, 1], [2, 1, 0]]) # Creating example array
myDict = {0: 19.5, 1: 34.1, 2: 10.9}
resultArray = ak.where(myArray == 0, myDict.get(0), 0)
for key in range(1,3):
resultArray = resultArray + ak.where(myArray == key, myDict.get(key), 0)

我的阵列：

<Array [[0, 1], [2, 1, 0]] type='2 * var * int64'>

结果数组：

<Array [[19.5, 34.1], [10.9, 34.1, 19.5]] type='2 * var * float64'>

Answer 1

当我在评论中提到np.searchsort是你应该寻找的地方时，我没有注意到myDict包含每个连续的整数作为键。拥有像这样的密集查找表将允许更快的算法，这在 Awkward Array 中也恰好更简单。

因此，假设从0到某个值的每个整数都有一个键myDict，您可以同样很好地将查找表表示为

>>> lookup = ak.Array([myDict[i] for i in range(len(myDict))])
>>> lookup
<Array [19.5, 34.1, 10.9] type='3 * float64'>

在0、1和2处选取值的问题变成了一个数组切片。(这个数组切片是数组长度 n 的 O(n) 算法，不像np.searchsorted是O(nlog n)。这就是使用稀疏查找键的成本。

然而，问题是myArray是嵌套的，而lookup不是。我们可以通过将其切成薄片来提供与myArray相同的深度lookup：

>>> multilookup = lookup[np.newaxis][np.zeros(len(myArray), np.int64)]
>>> multilookup
<Array [[19.5, 34.1, 10.9, ... 34.1, 10.9]] type='2 * 3 * float64'>
>>> multilookup.tolist()
[[19.5, 34.1, 10.9], [19.5, 34.1, 10.9]]

然后multilookup[myArray]正是您想要的：

>>> multilookup[myArray]
<Array [[19.5, 34.1], [10.9, 34.1, 19.5]] type='2 * var * float64'>

查找必须重复，因为myArray中的每个列表都使用整个lookup中的全局索引。如果创建multilookup所涉及的内存令人望而却步，则可以改为分解myArray以匹配它：

>>> flattened, num = ak.flatten(myArray), ak.num(myArray)
>>> flattened
<Array [0, 1, 2, 1, 0] type='5 * int64'>
>>> num
<Array [2, 3] type='2 * int64'>
>>> lookup[flattened]
<Array [19.5, 34.1, 10.9, 34.1, 19.5] type='5 * float64'>
>>> ak.unflatten(lookup[flattened], nums)
<Array [[19.5, 34.1], [10.9, 34.1, 19.5]] type='2 * var * float64'>

如果您的键从0到某个整数不密集，则必须使用np.searchsorted：

>>> keys = ak.Array(myDict.keys())
>>> values = ak.Array([myDict[key] for key in keys])
>>> keys
<Array [0, 1, 2] type='3 * int64'>
>>> values
<Array [19.5, 34.1, 10.9] type='3 * float64'>

在这种情况下，keys是微不足道的，因为它是密集的。使用np.searchsorted时，您必须将平面尴尬数组显式转换为 NumPy(目前;我们希望解决这个问题)。

>>> lookup_index = np.searchsorted(np.asarray(keys), np.asarray(flattened), side="left")
>>> lookup_index
array([0, 1, 2, 1, 0])

然后我们将其传递到琐碎的keys(在这种情况下，这不会改变它)，然后再将其传递给values。

>>> keys[lookup_index]
<Array [0, 1, 2, 1, 0] type='5 * int64'>
>>> values[keys[lookup_index]]
<Array [19.5, 34.1, 10.9, 34.1, 19.5] type='5 * float64'>
>>> ak.unflatten(values[keys[lookup_index]], num)
<Array [[19.5, 34.1], [10.9, 34.1, 19.5]] type='2 * var * float64'>

但是我在昨天的评论中胡说八道的是，你必须在扁平形式的myArray(flattened)上执行此操作，并在稍后重新引入结构ak.unflatten，如上所述。但也许我们应该np.searchsorted包装为ak.searchsorted，至少在第二个参数中识别一个完全结构化的笨拙数组。(它必须是非结构化的才能出现在第一个参数中。