如何使用ak.Array结合索引数组生成自定义掩码输出？

可以通过Awkward Array的矢量化操作完成需求，完全避免Python循环，适合大数据量场景：

实现步骤 #

1. 生成局部索引数组：创建与my_dummy_arr结构完全一致的数组，每个元素表示其在所属子数组中的位置。
2. 生成匹配掩码：逐子数组检查局部索引是否存在于my_indices对应的子数组中，得到布尔掩码。
3. 转换为整数结果：将布尔掩码转换为1/0的整数数组，得到最终结果。

完整代码 #

import awkward as ak

# 输入数组
my_indices = ak.Array([[0,1],[0],[1]])
my_dummy_arr = ak.Array([[1,1],[1,1],[1,1]])

# 1. 生成局部索引
local_idx = ak.local_index(my_dummy_arr)

# 2. 检查索引是否在对应my_indices子数组中（按子数组维度匹配）
mask = ak.isin(local_idx, my_indices, axis=1)

# 3. 转换为1/0整数数组
result = mask.astype(int)

print(result)
# 输出：[[1, 1], [1, 0], [0, 1]]

关键操作说明 #

• ak.local_index(my_dummy_arr)：生成结构与输入完全对齐的局部索引，例如[[0,1],[0,1],[0,1]]，确保后续匹配的结构一致性。
• ak.isin(..., axis=1)：指定按子数组（轴1）进行元素匹配，保证每个子数组的索引只和对应的my_indices子数组对比，而非全局匹配。
• astype(int)：直接将布尔值True/False转换为1/0，是Awkward内部的矢量化操作，性能远高于Python循环。

内容的提问来源于stack exchange，提问作者Crucio