P5568. NPU梯度任务均衡调度

场次：2026年9月23日 · 来源：CodeFun2000 P5568

题目描述

在大模型分布式训练场景下，要把若干梯度更新任务分配到多张 NPU 上。目标是让各卡的「梯度计算总耗时」尽量均衡（即最小化「最大卡总耗时」），再按学习率衰减规则汇总总训练成本。

规则如下：

1. 任务分配：将 $grad\_times$ 划分成 $k$ 个子集（一张 NPU 对应一个子集）。每个任务必须分到恰好一个子集；在所有划分中，取「子集总和的最大值」最小的那些划分。
2. 成本计算：记 $max\_time$ 为各卡总耗时的最大值，$sum\_time\_j$ 为第 $j$ 张卡的总耗时。则
   • 第 $j$ 张卡学习率：$lr\_j = base\_lr \times decay\_rate^{max\_time - sum\_time\_j}$；
   • 第 $j$ 张卡训练成本：$sum\_time\_j \times lr\_j$；
   • 总训练成本：所有卡成本之和。

输入格式

一行字符串，用 ; 分成四段：梯度耗时列表;NPU数量;基础学习率;衰减系数。例如：$8,5,4,3,3,2,1;3;0.1;0.8$。

1. 梯度耗时列表（$grad\_times$）：逗号分隔的一维数组，每个元素为单个梯度更新任务的耗时（非负整数）；
2. $NPU$ 数量（$k$）：正整数，满足 $1 \le k \le len(grad\_times)$；
3. 基础学习率（$base\_lr$）：浮点数；
4. 衰减系数（$decay\_rate$）：浮点数，满足 $0 < decay\_rate < 1$。

输出格式

一行字符串：最小最大总耗时,总训练成本。例如：$9,2.4400$。

• 最小最大总耗时：整数；
• 总训练成本：保留 $4$ 位小数的浮点数。

示例

示例 1

输入：

8,5,4,3,3,2,1;3;0.1;0.8

输出：

9,2.4400

说明：

解析得 $grad\_times=[8,5,4,3,3,2,1]$，$k=3$，$base\_lr=0.1$，$decay\_rate=0.8$。

总耗时 $8+5+4+3+3+2+1=26$，理想均分约 $26/3\approx 8.67$，且最大单任务为 $8$，故最小可能的最大子集和至少为 $9$。

一组达到该下界的划分：

• NPU1：$8+1=9$
• NPU2：$5+4=9$
• NPU3：$3+3+2=8$

负载为 $9,9,8$，$max\_time=9$。

成本：

• 负载 $9$：$9\times 0.1\times 0.8^{0}=0.9$（两张）
• 负载 $8$：$8\times 0.1\times 0.8^{1}=0.64$

总成本 $0.9+0.9+0.64=2.44$，输出 $9,2.4400$。

示例 2

输入：

6;1;0.2;0.9

输出：

6,1.2000

说明：

仅 $1$ 张 NPU、一个任务，故 $max\_time=6$，$sum\_time=6$。

$lr=0.2\times 0.9^{0}=0.2$，总成本 $6\times 0.2=1.2$，输出 $6,1.2000$。

---

解题思路

第一步：理解问题本质

题目有两个先后阶段，顺序不能颠倒：

1. 先决定「怎么分」：把 $grad\_times$ 分成 $k$ 组（一组对应一张 NPU），要求每组耗时之和的最大值尽可能小，记这个最小值为 $max\_time$；
2. 再算「花多少钱」：用第一阶段得到的各卡负载 $sum\_time_j$ 代入

$$lr_j = base\_lr\times decay\_rate^{\,max\_time-sum\_time_j},\qquad cost=\sum_{j=1}^{k} sum\_time_j\times lr_j$$

需要注意公式里的指数是 $max\_time-sum\_time_j$：负载越小的卡，指数越大，而 $0<decay\_rate<1$，指数越大结果越小，所以负载小的卡学习率反而更低、成本更小。也就是说「总成本最低」和「负载均衡」是两件不同的事，题目要求的顺序是先均衡（最小化最大负载），再按公式算钱。

第一步本质上是经典的「多机调度」问题：把若干任务分给 $k$ 台相同的机器，最小化最大完工时间。它是 NP-hard 的，没有多项式时间的精确算法，因此题目的任务数不会太大，可以放心用枚举 + 剪枝来精确求解。

第二步：暴力解法

每个任务都有 $k$ 种归宿，$n$ 个任务一共 $k^n$ 种分配。直接枚举所有分配，记录每张卡的负载，找出「最大负载最小」的方案：

best_max = float("inf")
best_loads = None

for assign in product(range(k), repeat=n):   # 每个任务放第几张卡
    loads = [0] * k
    for t, card in zip(times, assign):
        loads[card] += t
    cur_max = max(loads)
    if cur_max < best_max:                   # 只关心最大负载
        best_max = cur_max
        best_loads = loads[:]

# 拿到最优负载后再套衰减公式
cost = sum(s * base_lr * decay ** (best_max - s) for s in best_loads)

$n=10$、$k=3$ 时是 $59049$ 种，勉强能跑；$n=20$、$k=3$ 就是 $3.5\times10^9$，直接超时。而且它把「所有叶子」都走了一遍，其实我们在找到一个不错的方案之后，就可以把「已经比它还差」的分支整片砍掉。

第三步：优化解法——把「最小化」变成「判定」再二分

与其直接求最小值，不如先回答一个更容易的问题：

给定上限 $cap$，能不能把任务分下去，使得每张卡的耗时都不超过 $cap$？

这个判定问题有两个好处：

1. 单调性：如果 $cap$ 能做到，那么比它大的 $cap'$ 也一定能做到（原来的方案照样合法）。于是「能做到」的 $cap$ 构成一段后缀 $[\,max\_time,\ +\infty)$，可以二分。
2. 判定可以「提前收工」：只要找到一条可行路径就能返回 True，不需要枚举所有方案。配上「放不下就回头」的剪枝，搜索树会被砍得非常小。

二分的范围：

• 下界：最大的单个任务必须整块放进某张卡，所以 $cap\ge\max(grad\_times)$；同时 $k$ 张卡平均分摊总和 $S$，所以 $cap\ge\lceil S/k\rceil$。两者取较大值。
• 上界：把所有任务都堆到一张卡上，最大负载就是 $S$，一定可行。

二分大约只需要 $\log_2 S$ 次判定，就能得到最小的 $max\_time$。

第四步：最优解法——带回溯剪枝的判定 + 复搜一次取负载

判定函数就是一次 DFS：第 $i$ 层决策「第 $i$ 个任务放哪张卡」。

dfs(i):
    如果 i == n：找到可行方案，返回 True
    对每张卡 j：
        如果 loads[j] + times[i] > cap：放不下，跳过
        loads[j] += times[i]        （做选择）
        如果 dfs(i+1) 成功：返回 True
        loads[j] -= times[i]        （撤销选择，回溯）
    返回 False

在上面这个骨架上，加三层剪枝：

1. 上限剪枝：loads[j] + times[i] > cap 直接跳过。这是最主要的剪枝，二分出的 $cap$ 很紧，大部分分支在这里就被砍掉。
2. 对称剪枝：如果两张卡当前的负载完全相同，那么把任务放进哪一张得到的局面完全一样。所以同一层只需要尝试「不同的负载值」各一次，用一个小集合 tried 记录已经试过的负载即可。
3. 失败记忆：局面可以用 (当前处理到第几个任务, 排序后的负载列表) 唯一表示——卡之间可以互换，所以负载要先排序。同一个局面如果已经证明无解，就没必要再搜一遍。

得到最小的 $max\_time$ 之后，用同样的判定函数（$cap=max\_time$）再搜一次，这次把成功时的负载数组留下来，就得到了一组达到最小值的 $sum\_time_j$。最后按公式累加成本、格式化输出。

关于并列最优：如果存在多组划分都能达到同样的 $max\_time$，题面没有规定进一步的比较规则，代码会取搜索顺序遇到的第一组负载。当测试数据的最优划分唯一时，这样得到的结果就是唯一正确的那个。

---

完整代码实现

"""
P5568. NPU梯度任务均衡调度 - 二分答案 + 回溯搜索

输入格式：一行字符串，用 ; 分成四段：
    grad_times;NPU数量;base_lr;decay_rate

输出格式：一行 "最小最大总耗时,总训练成本（4 位小数）"

复杂度：二分 O(log S) 次判定，每次判定是带剪枝的回溯搜索，
        最坏 O(k^n)，实际数据规模下剪枝后很快。
"""

import sys


def search(times, k, cap):
    """在「每张卡耗时都不超过 cap」的前提下找一个分配方案。

    找到时返回各卡负载列表（长度 k），找不到返回 None。
    """
    n = len(times)
    loads = [0] * k

    # 已经证明无解的局面：卡可以互换，所以负载要先排序再去重
    failed = set()

    def dfs(i):
        if i == n:            # 所有任务分配完毕，这一组负载满足上限
            return True

        state = (i, tuple(sorted(loads)))
        if state in failed:
            return False

        t = times[i]
        tried = set()         # 本层已经尝试过的负载值

        for j in range(k):
            cur = loads[j]
            # 负载相同的卡完全对称，只需要试一次
            if cur in tried:
                continue
            # 放进去就超上限，这条分支不可能有合法方案
            if cur + t > cap:
                continue

            tried.add(cur)
            loads[j] = cur + t
            if dfs(i + 1):
                return True
            loads[j] = cur    # 回溯：撤销这次分配
        failed.add(state)
        return False

    return loads if dfs(0) else None


def solve():
    text = sys.stdin.buffer.read().decode().strip()
    if not text:
        return

    parts = text.split(";")
    if len(parts) < 4:
        return

    times = [int(x) for x in parts[0].split(",") if x.strip()]
    k = int(parts[1])
    base_lr = float(parts[2])
    decay = float(parts[3])

    n = len(times)
    if k > n:                 # 题面保证 k <= n，这里兜一层
        k = n

    total = sum(times)

    # 二分最小的 max_time
    low = max(max(times), (total + k - 1) // k)
    high = total

    while low < high:
        mid = (low + high) // 2
        if search(times, k, mid) is not None:
            high = mid
        else:
            low = mid + 1

    max_time = low

    # 用最小的上限再搜一次，拿到一组具体的负载分配
    loads = search(times, k, max_time)

    # 按衰减公式累计每张卡的成本
    cost = 0.0
    for s in loads:
        cost += s * base_lr * (decay ** (max_time - s))

    print(f"{max_time},{cost:.4f}")


if __name__ == "__main__":
    solve()

---

示例推演

样例 1

输入：

8,5,4,3,3,2,1;3;0.1;0.8

解析得 $grad\_times=[8,5,4,3,3,2,1]$，$k=3$，$base\_lr=0.1$，$decay\_rate=0.8$，总和 $S=26$。

第一步：确定二分的下界。

• 最大单任务 $=8$；
• 平均值 $\lceil 26/3\rceil=\lceil 8.67\rceil=9$；
• 下界取 $\max(8,9)=9$，上界取 $26$。

第二步：验证 $cap=9$ 可行、$cap=8$ 不可行。

• $cap=8$：三张卡最多只能装 $3\times8=24<26$，装不下全部任务，不可行；
• $cap=9$：存在方案，例如

NPU	任务	负载
1	8 + 1	9
2	5 + 4	9
3	3 + 3 + 2	8

所以 $max\_time=9$（已经等于下界，二分不会再往下走）。

第三步：按搜索顺序拿到负载。

第 $1$ 个任务 $8$ 放进卡 1（$8$）；$5$ 放进卡 2（$8,5$）；$4$ 放卡 2（$8,9$）；$3$ 放卡 3（$8,9,3$）；$3$ 放卡 3（$8,9,6$）；$2$ 放卡 3（$8,9,8$）；$1$ 放卡 1（$9,9,8$）。 每一步都没有超过 $9$，最终负载为 $(9,9,8)$，与上面的方案一致。

第四步：算成本。

• 负载 $9$：$9\times0.1\times0.8^{\,9-9}=0.9$，有两张；
• 负载 $8$：$8\times0.1\times0.8^{\,9-8}=0.8\times0.8=0.64$。

$$cost=0.9+0.9+0.64=2.44$$

输出：

9,2.4400

样例 2

输入：

6;1;0.2;0.9

$n=1$、$k=1$：只有一个任务，只能放唯一一张卡。

• 二分下界 $\max(6,\lceil6/1\rceil)=6$，上界 $6$，所以 $max\_time=6$；
• 负载就是 $(6)$，指数 $6-6=0$；
• 成本 $6\times0.2\times0.9^0=1.2$。

输出：

6,1.2000

---

复杂度分析

解法	时间复杂度	空间复杂度	说明
枚举全部 $k^n$ 种分配	$O(k^n\cdot k)$	$O(n+k)$	直接得答案，$n$ 稍大就超时
二分答案 + 朴素回溯判定	$O(\log S\cdot k^n)$	$O(n+k)$	判定能提前返回，剪枝更强
二分答案 + 对称去重 + 失败记忆	$O(\log S\cdot M)$	$O(n+k+M)$	本题采用，$M$ 为实际搜到的局面数，远小于 $k^n$

其中 $S=\sum grad\_times$ 是总耗时，也决定了二分的上界。二分的次数是 $O(\log S)$，而每次判定一旦找到可行方案就立刻返回，配合「放不下就剪枝」与对称去重，实际搜索量很小。

---

易错点总结

1. 输出格式是「整数,四位小数」

最小最大耗时输出整数，成本固定保留 4 位小数，两者用逗号连接、中间没有空格：9,2.4400。用 f"{max_time},{cost:.4f}" 最直接；round(cost, 4) 不会补零，2.4 会变成 2.4 而不是 2.4000，输出直接对不上。

2. 指数是 $max\_time-sum\_time_j$，不要写反

负载恰好等于 $max\_time$ 的卡，指数为 $0$，$lr=base\_lr$；负载越轻，指数越大，学习率被衰减得越厉害。写成 $sum\_time_j-max\_time$ 会让负载轻的卡拿到更大的学习率，成本整体飙升。

3. 下界必须考虑最大单任务

只取 $\lceil S/k\rceil$ 会低估：样例 1 如果总和是 $26$、但某个任务是 $20$，那么 $cap$ 至少要 $20$。下界取 max(max(times), ceil(total/k)) 既正确又能少做几次判定。

4. 回溯要真正撤销

loads[j] += t 之后必须 loads[j] -= t，否则后面的兄弟分支会继承上一个分支的负载，搜索结果完全错乱。这是回溯最容易漏的一步。

5. 对称去重不会漏解

两张卡当前负载相同，把任务放进哪一张，得到的「负载集合」是同一个，所以只试一次不会丢解。反过来，如果不做这一步，$k$ 大的时候会白白多搜 $k!$ 量级的重复分支。

6. 失败记忆的键要先排序

(i, tuple(loads)) 这种写法会把「卡 1 是 5、卡 2 是 3」和「卡 1 是 3、卡 2 是 5」当成两个不同局面，因为它们对答案完全等价，必须 tuple(sorted(loads)) 才能去重。

7. 允许空卡与 0 耗时任务

$k\le n$，但搜索过程中允许某张卡暂时没有任务（负载 $0$），它的成本贡献是 $0\times lr=0$，不影响答案。任务耗时为 $0$ 也是合法的，此时下界可能为 $0$，low = high = 0 会直接跳过二分。

8. 并列最优时的负载选择

最小 $max\_time$ 可能对应多组负载，而不同负载代入衰减公式会得到不同的成本。题面没有给出进一步的规则，代码取搜索顺序的第一组。如果评测数据的最优划分唯一（通常如此），这一点不会带来差异。

---

扩展思考

• 多机调度是 NP-hard：本题的精确解法只能在 $n$ 较小时使用。规模大时可以用 LPT 贪心（把任务从大到小排序，每次放到当前负载最小的卡上），最坏情况下最大负载不超过最优值的 $\frac43-\frac{1}{3k}$ 倍。
• 另一种精确解法：按子集做 DP，把「哪些任务已经分好」压成 bitmask，再枚举最后一个子集作为一张卡的任务，复杂度 $O(3^n)$，适合 $n\le 15$ 且 $k$ 不确定的场景。
• 如果限制「任务必须连续分配」：问题就退化成「分割数组的最大值」（LeetCode 410），可以直接用二分 + 贪心判定，不需要枚举子集。
• 成本的形状：把每张卡的成本看成关于负载的函数 $f(s)=s\cdot base\_lr\cdot decay^{\,max\_time-s}$，它在 $s\in[0,max\_time]$ 上单调递增、而且是凸函数。所以在总耗时固定时，负载越均衡，总成本越低；「把任务尽量摊平」既能压小 $max\_time$，也能压小成本，两个目标大方向一致。不过成本还取决于每张卡具体的负载分布，题目仍明确以 $max\_time$ 为第一目标。
• 浮点输出：decay ** 整数 与 base_lr 都是二进制浮点数，末位可能存在极小的误差。按题目要求保留 4 位小数，正常范围内不会出现四舍五入到错误值的情况；如果担心边界，可以改用 decimal 提高精度。

---

相关题目

• 激活检查点放置
• 回溯
• 递归