Numpy 3D数组中子数组存在性高效检测与重复项处理

本文探讨了如何在numpy 3d数组中高效检测一个子数组是否存在于另一个可能更短且包含重复项的3d数组中，并返回布尔结果。文章提供了两种主要方法：通过字符串转换结合`np.in1d`进行一维比较，以及利用numpy的广播机制进行直接的元素级比较，并详细解释了它们的实现原理、适用场景及潜在的性能考量。

在数据处理和科学计算中，我们经常需要比较大型多维数组中的元素。一个常见的场景是，给定一个包含多个三维向量（或称为2D子数组）的3D NumPy数组source，我们需要判断其中每个向量是否存在于另一个可能更短且包含重复向量的3D数组values中。例如，我们有以下两个数组：

import numpy as np

source = np.array([[[0,0,0],[0,0,1],[0,1,0],[1,0,0],[1,0,1],[1,1,0],[1,1,1]]])
values = np.array([[[0,1,0],[1,0,0],[1,1,1],[1,1,1],[0,1,0]]])

我们的目标是得到一个布尔数组，其长度与source数组在Axis1上的长度相同，表示source中每个2D子数组是否在values中出现。对于上述示例，期望的输出是[False, False, True, True, False, False, True]。

直接使用np.isin(source, values).all(axis=2)通常无法达到预期效果，因为它会在元素层面进行比较，而不是在完整的2D子数组层面。np.in1d虽然适用于一维数组，但直接应用于3D数组时需要巧妙的转换。

方法一：字符串转换与np.in1d

这种方法的核心思想是将每个2D子数组（例如[0,0,0]）转换为一个唯一的字符串表示，从而将3D数组的比较问题转化为一维字符串数组的比较问题，进而可以利用np.in1d的强大功能。

result_string_conversion = np.in1d(np.apply_along_axis(''.join, 2, source.astype(str)), 
                                     np.apply_along_axis(''.join, 2, values.astype(str)))
print(result_string_conversion)
# 输出: [False False  True  True False False  True]

实现原理详解：

1. source.astype(str) 和 values.astype(str)：首先，将source和values数组中的所有数值元素转换为字符串类型。这是因为''.join函数只能连接字符串。
2. np.apply_along_axis(''.join, 2, ...)：这是关键一步。
   • ' '.join是一个字符串连接函数，它将一个可迭代对象中的所有字符串连接成一个单一的字符串。
   • axis=2指定了操作的轴。对于source数组（形状为(1, 7, 3)），axis=2指的是最内层的维度，即[0,0,0]中的0,0,0。apply_along_axis会沿着这个轴对每个2D子数组（如[0,0,0]）应用''.join函数。
   • 例如，[0,0,0]会转换为字符串"000"，[0,0,1]转换为"001"，依此类推。
   • 这样，原始的3D数组就被有效地“扁平化”为一维的字符串数组，其中每个字符串代表一个原始的2D子数组。
3. np.in1d(array1, array2)：最后，np.in1d函数用于判断array1中的每个元素是否存在于array2中，返回一个布尔数组。在这里，array1是source转换后的字符串数组，array2是values转换后的字符串数组。

注意事项：

• 这种方法在处理数值数组时需要额外的字符串转换步骤，可能会引入一定的性能开销。
• 对于非常大的数组，生成大量的中间字符串也可能占用较多内存。
• 它依赖于将每个子数组转换为唯一字符串，如果子数组元素的值范围很大，或者包含浮点数精度问题，需要确保字符串转换能准确表示唯一性。

方法二：基于广播的直接比较

第二种方法利用NumPy的广播机制进行直接的元素级比较，然后通过聚合操作来判断子数组的匹配情况。这种方法在某些情况下可能更直接，但需要对NumPy的广播规则有较好的理解。

STORYD

帮你写出让领导满意的精美文稿

下载

# 确保source和values的维度兼容，如果source是(1, N, M)而values是(1, K, M)，
# 需要调整为 (N, M) 和 (K, M) 进行比较，或者在比较前进行适当的维度扩展。
# 考虑到示例数据，source[0] 和 values[0] 才是实际要比较的二维数组
# source_2d = source[0] # 形状 (7, 3)
# values_2d = values[0] # 形状 (5, 3)

# 原始答案中的方法假设了source和values的维度结构，我们将其适配到 (N, M) 和 (K, M) 的比较
# 为了保持与原始答案一致，我们使用其提供的代码，它隐式处理了第一维
result_broadcast_comparison = (source.transpose(1,0,2) == values).all(2).any(1)
print(result_broadcast_conversion)
# 输出: [False False  True  True False False  True]

实现原理详解：

为了更好地理解，我们先假设我们要比较的是source[0]（形状 (7, 3)）和 values[0]（形状 (5, 3)）。

1. 维度调整与广播准备：
   • source.transpose(1,0,2)：这一步是关键，它将source的维度从(1, N, M)（例如(1, 7, 3)）转换为(N, 1, M)（例如(7, 1, 3)）。
     • 为什么要这样做？当它与values（形状(1, K, M)，例如(1, 5, 3)）进行比较时，NumPy的广播规则会生效。
     • source现在是(7, 1, 3)，values是(1, 5, 3)。
     • 广播时，NumPy会扩展维度为1的轴，使其与另一个数组的对应轴长度相同。
     • 结果，source会被扩展为(7, 5, 3)，values也会被扩展为(7, 5, 3)。
     • 这意味着source的每个子数组（例如[0,0,0]）都会与values中的所有子数组（例如[0,1,0], [1,0,0], [1,1,1], [1,1,1], [0,1,0]）进行逐一比较。
2. == values：执行元素级的相等性比较。这将产生一个布尔数组，形状为(7, 5, 3)。output[i, j, k]为True如果source[i, :, k]等于values[j, :, k]。
3. .all(2)：沿着最后一个轴（Axis 2，即每个子数组内部的元素）进行all操作。
   • 对于形状为(7, 5, 3)的布尔数组，all(2)会检查source的第i个子数组是否与values的第j个子数组的所有元素都匹配。
   • 这将把数组的形状从(7, 5, 3)缩减为(7, 5)。output[i, j]为True表示source的第i个子数组与values的第j个子数组完全匹配。
4. .any(1)：沿着倒数第二个轴（Axis 1，即values中的所有子数组）进行any操作。
   • 对于形状为(7, 5)的布尔数组，any(1)会检查source的第i个子数组是否与values中的任何一个子数组完全匹配。
   • 这将把数组的形状从(7, 5)缩减为(7,)，得到最终的布尔结果数组。

注意事项：

• 这种方法在中间步骤会生成一个非常大的临时数组（形状为(N, K, M)），其内存消耗可能非常高，尤其是当N和K都很大时。例如，如果source有1000个子数组，values有1000个子数组，那么中间数组将有1000 1000 M个元素。
• 尽管内存消耗大，但纯NumPy操作通常在CPU层面经过高度优化，对于中等大小的数组，其计算速度可能非常快。

总结

在NumPy中检测3D数组中子数组的存在性，尤其是在存在重复项的情况下，需要巧妙地运用NumPy的特性。

• 字符串转换与np.in1d 方法通过将复杂的3D比较问题转化为简单的1D字符串比较，提供了一种通用且内存效率相对较高（尽管有字符串转换开销）的解决方案。它适用于各种数据类型，只要能转换为唯一的字符串表示即可。
• 基于广播的直接比较 方法则利用了NumPy强大的广播机制，直接在数值层面进行比较。它在计算上可能非常高效，但以潜在的巨大内存消耗为代价，因此更适合于source和values数组在相关维度上不是特别大的场景。

选择哪种方法取决于您的具体需求：如果内存是关键限制，或者子数组元素类型复杂，字符串转换可能是更好的选择。如果计算性能至关重要且内存允许，广播方法可能提供更快的执行速度。理解这两种方法的内在机制，有助于在实际应用中做出明智的决策。