通过另一个 JaggedArray 有效地排序和过滤

Question

我有一个 JaggedArray (awkward.array.jagged.JaggedArray)，其中包含指向另一个 JaggedArray 中位置的索引。两个数组的长度相同，但 JaggedArrays 包含的每个 numpy.ndarrays 的长度可以不同。我想使用第一个数组的索引对第二个数组进行排序，同时从第二个数组中删除未从第一个数组索引的元素。第一个数组还可以包含 -1 的值（如果需要，也可以用 None 替换，但目前不是这种情况），这意味着第二个数组中没有匹配项。在这种情况下，第一个数组中的相应位置应设置为默认值（例如 0）。

这是一个实际的例子，以及我目前是如何解决这个问题的：

import uproot
import numpy as np
import awkward

def good_index(my_indices, my_values):
    my_list = []
    for index in my_indices:
        if index > -1:
            my_list.append(my_values[index])
        else:
            my_list.append(0)
    return my_list

indices = awkward.fromiter([[0, -1], [3,1,-1], [-1,0,-1]])
values = awkward.fromiter([[1.1, 1.2, 1.3], [2.1,2.2,2.3,2.4], [3.1]])

new_map = awkward.fromiter(map(good_index, indices, values))

结果 new_map 是：[[1.1 0.0] [2.4 2.2 0.0] [0.0 3.1 0.0]].

是否有更多 efficient/faster 方法来实现此目的？我在想可以使用 numpy 功能，例如 numpy.where，但由于 ndarrays 的长度不同，至少对于我尝试的方式来说，这失败了。

Answer 1

如果保证values中的所有子数组都是非空的（这样用-1returns索引最后一个子元素，不是错误），那么你可以这样做：

>>> almost = values[indices]       # almost what you want; uses -1 as a real index
>>> almost.content = awkward.MaskedArray(indices.content < 0, almost.content)
>>> almost.fillna(0.0)
<JaggedArray [[1.1 0.0] [2.4 2.2 0.0] [0.0 3.1 0.0]] at 0x7fe54c713c88>

最后一步是可选的，因为没有它，缺少的元素是 None，而不是 0.0。

如果 values 中的某些子数组为空，您可以 pad 它们以确保它们至少有一个子元素。所有原始子元素的索引方式与之前相同，因为 pad 仅在需要时增加长度。

>>> values = awkward.fromiter([[1.1, 1.2, 1.3], [], [2.1, 2.2, 2.3, 2.4], [], [3.1]])
>>> values.pad(1)
<JaggedArray [[1.1 1.2 1.3] [None] [2.1 2.2 2.3 2.4] [None] [3.1]] at 0x7fe54c713978>

通过另一个 JaggedArray 有效地排序和过滤

Efficiently sorting and filtering a JaggedArray by another one

awkward-array