“组队”刷题实践 | 豆包MarsCode AI刷题

66 阅读7分钟

今天刷一刷豆包MarsCode AI刷题中的“组队”

{EA4868CA-84E4-4A56-9726-5D868E08F790}.png

题目内容

题目要求我们将公司内的 n 个人分成至少 k 个非空小组,每个人只能属于一个组。每个人的能力值和性格值分别用 ai 和 bi 表示。两个人的差别值定义为能力值之差的绝对值加上性格值之差的绝对值。

分组的要求是,如果两个人的差别值不超过某个上限 L,那么他们必须在同一个组内。我们的目标是确定在至少有 k 个组的前提下,差别上限 L 的最大值是多少。

思路讲解

这个问题实际上是一个典型的二分搜索与图论结合的问题。主要会用到 【二分查找】【并查集】

为什么使用二分和并查集:

  • 二分:因为我们要求的是差别上限 L 的最大值,这是一个有序的搜索问题,二分搜索可以高效地缩小搜索范围,提高解决问题的效率。
  • 并查集:用于处理动态连通性问题,即判断图中的节点是否属于同一个连通分量。在这个问题中,我们需要判断在给定的差别上限下,能否将所有人分成至少 k 个组,并查集可以高效地完成这一任务。 通过结合二分搜索和并查集,我们可以高效地找到差别上限 L 的最大值。

二分搜索的思路:

  1. 确定搜索范围:由于我们要找的是差别上限 L 的最大值,所以我们可以将 L 的可能值作为搜索的范围。最小值可以是 0,最大值可以是所有人能力值和性格值之差的最大可能值。
  2. 二分搜索:在上述范围内,我们可以使用二分搜索来确定最大的 L 值。对于每一个中间值 mid,我们检查是否可以将所有人分成至少 k 个组,并且组内任意两个人的差别值不超过 mid。
  3. 调整搜索范围:如果当前 mid 值可以满足条件,那么说明 L 的最大值可能更大,我们应该将搜索范围的上界上调;如果当前 mid 值不满足条件,那么说明 L 的最大值必须更小,我们应该将搜索范围的下界下调。

并查集的思路:

  1. 图的构建:将每个人看作图中的一个节点,如果两个人的差别值不超过当前的 mid 值,那么这两个节点之间就有一条边。
  2. 查找连通分量:使用并查集来查找图中的连通分量。每个连通分量代表一个组,组内任意两个人的差别值不超过 mid。
  3. 判断是否满足条件:如果连通分量的数量少于 k 个,那么当前的 mid 值不满足条件;如果连通分量的数量大于或等于 k 个,那么当前的 mid 值满足条件。

具体思路

  1. 初始化:首先,我们初始化一个并查集,用于管理每个人的分组情况。
  2. 二分查找:我们需要找到最大的差别上限 L,这可以通过二分查找来实现。我们将 L 的可能值范围设置为 [0, max(max(a) - min(a), max(b) - min(b))]。
  3. 分组检查:在二分查找的每一步,我们使用当前的 L 值来检查是否可以将所有人分成至少 k 个组。具体做法是遍历所有人,如果两个人的差别值小于等于 L,则将他们合并到同一个组中。
  4. 判断分组数量:在当前的 L 值下,如果最终的分组数量大于等于 k,则说明 L 还可以更大,我们将 l 更新为 mid;否则,说明 L 太大了,需要减小,我们将 r 更新为 mid - 1。
  5. 返回结果:当 l 和 r 相遇时,我们就找到了最大的差别上限 L。

概念讲解

并查集是什么

并查集(Union-Find)是一种数据结构,主要用于处理一些不交集的合并及查询问题。它支持两种操作:

  1. 查找(Find):确定某个元素属于哪个子集。它可以被用来确定两个元素是否属于同一个子集。
  2. 合并(Union):将两个子集合并成一个集合。 并查集通常使用一个数组来表示,数组的索引代表元素,而数组中的值代表该元素的父亲节点。

image.png

如何实现

下面给出代码模板,实现并查集(Union-Find)数据结构的两个核心操作:findunion

find 操作

find 操作用于查找元素 i 所在的集合的根节点,这个根节点也被称为集合的代表元素。如果元素 i 已经在集合的根节点上,那么直接返回 i;否则,递归地查找 parent[i] 的根节点,并在返回的过程中进行路径压缩,即直接将 i 的父节点指向其根节点,这样可以减少后续 find 操作的递归深度。

def find(parent: List[int], i: int) -> int:
    if parent[i] == i:  # 如果当前节点是根节点,直接返回
        return i
    parent[i] = find(parent, parent[i])  # 递归查找根节点,并进行路径压缩
    return parent[i]

union 操作

union 操作用于将两个元素 xy 所在的集合合并成同一个集合。首先,分别找到 xy 的根节点 xrootyroot。如果 xy 已经在同一个集合中(即根节点相同),则不需要任何操作。否则,根据两个根节点的秩(rank)来决定合并方向:

  • 如果 xroot 的秩小于 yroot 的秩,将 xroot 指向 yroot
  • 如果 xroot 的秩大于 yroot 的秩,将 yroot 指向 xroot
  • 如果两者秩相同,则将 yroot 指向 xroot 并将 xroot 的秩加一。

这样做的目的是为了保持树的平衡,减少树的高度,从而提高 find 操作的效率。

def union(parent: List[int], rank: List[int], x: int, y: int) -> None:
    xroot = find(parent, x)  # 找到x的根节点
    yroot = find(parent, y)  # 找到y的根节点
    
    if xroot != yroot:  # 如果x和y不在同一个集合中
        if rank[xroot] < rank[yroot]:  # 比较秩,决定合并方向
            parent[xroot] = yroot
        elif rank[xroot] > rank[yroot]:
            parent[yroot] = xroot
        else:
            parent[yroot] = xroot  # 秩相同,任意选择一个作为根节点
            rank[xroot] += 1  # 增加根节点的秩

总结

并查集是一种非常高效的数据结构,用于处理一些不交集的合并及查询问题。通过 findunion 操作,我们可以快速地确定元素是否在同一集合中,以及合并两个集合。在上述代码中,find 操作通过路径压缩优化了查找效率,而 union 操作通过秩优化了树的平衡性,两者共同保证了并查集操作的高效性。

完整实现

以下是本题的完整的代码实现:

from typing import List
# 查找操作,带路径压缩
def find(parent: List[int], i: int) -> int:
    if parent[i] == i:
        return i
    parent[i] = find(parent, parent[i])
    return parent[i]
# 合并操作,按秩合并
def union(parent: List[int], rank: List[int], x: int, y: int) -> None:
    xroot = find(parent, x)
    yroot = find(parent, y)
    
    if xroot != yroot:
        if rank[xroot] < rank[yroot]:
            parent[xroot] = yroot
        elif rank[xroot] > rank[yroot]:
            parent[yroot] = xroot
        else:
            parent[yroot] = xroot
            rank[xroot] += 1
# 主函数,使用二分查找和并查集来找到最大的差别上限 L
def solution(n: int, k: int, a: List[int], b: List[int]) -> int:
    l, r = 0, max(max(a) - min(a), max(b) - min(b))
    
    while l < r:
        mid = (l + r + 1) // 2
        parent = list(range(n))
        rank = [0] * n
        
        for i in range(n):
            for j in range(i + 1, n):
                if abs(a[i] - a[j]) + abs(b[i] - b[j]) <= mid:
                    union(parent, rank, i, j)
        
        groups = len(set(find(parent, i) for i in range(n)))
        
        if groups >= k:
            l = mid
        else:
            r = mid - 1
            
    return l
# 测试样例
if __name__ == '__main__':
    print(solution(3, 2, [1, 9, 3], [2, 7, 8]) == 

{D0E17BCD-A662-4D4B-A37A-69FE50E28331}.png