如何判断输入数组的元素是否来自同一分组？

给定分组列表：

a = [[658,659,660],[1217,1218,1219],[1691,1692,1693],[1694,1695,1696]]

要求实现一个判断逻辑：仅当输入数组的所有元素都来自同一个嵌套分组时返回True，元素跨分组则返回False。

示例输入输出 #

• [658] → True
• [658,659] → True
• [658,659,660] → True
• [1217] → True
• [1217,1218] → True
• [1217,1218,1219] → True
• [1217,658] → False
• [1217,1218,658] → False
• [658,659,1217] → False

方法一：预构建映射（高效查询） #

先建立元素到所属分组的字典，后续查询时只需检查所有元素是否属于同一分组，适合频繁查询的场景：

# 预生成元素-分组映射
group_map = {}
for group in a:
    for num in group:
        group_map[num] = group

def is_same_group(input_arr):
    if not input_arr:
        return True  # 空数组的返回值可根据实际需求调整
    target_group = group_map.get(input_arr[0])
    # 检查所有元素是否都在目标分组内
    return all(num in target_group for num in input_arr)

测试示例：

print(is_same_group([658]))          # True
print(is_same_group([1217,658]))     # False
print(is_same_group([1694,1695]))    # True

方法二：实时遍历分组（无需预构建） #

如果分组列表规模较小，也可以直接遍历分组统计元素所属的分组数量，数量为1则返回True：

def is_same_group(input_arr):
    if not input_arr:
        return True
    # 用集合存储元素所属的分组（转元组实现哈希）
    unique_groups = set()
    for num in input_arr:
        for group in a:
            if num in group:
                unique_groups.add(tuple(group))
                break
    return len(unique_groups) == 1

内容的提问来源于stack exchange，提问作者Hemanth Hema Surya