P5468. 流水线并行的最小峰值负载

场次：2026年9月18日 · 来源：CodeFun2000 P5468

题目描述

要把一个神经网络的 $N$ 个层切到 $K$ 块 GPU 上做流水线并行。切完后每一段叫做一个 Stage，训练时按 Stage 顺序接力计算。

一个 Stage 的负载等于它包含的各层前向传播时间之和。若 Stage $s$ 覆盖连续层 $l_1,\dots,l_t$，则

$Load_s = f[l_1] + \cdots + f[l_t]$

其中 $f[i]$ 是第 $i$ 层的前向时间。流水线由最慢的 Stage 卡住，因此要最小化所有 Stage 负载的最大值。

切分必须同时满足：

1. 每个 Stage 内的层编号连续
2. 每一层恰好属于一个 Stage
3. 每个 Stage 至少含一层

输入格式

第一行两个整数 $N$、$K$，依次为层数和 Stage 个数。

第二行 $N$ 个整数 $f[0]\ f[1]\ \dots\ f[N-1]$，表示各层前向时间（毫秒）。

输出格式

输出一行一个整数，即最小化后的最大 Stage 负载。

数据范围

• $1 \le K \le N \le 5000$
• $1 \le f[i] \le 1000$

示例

示例 1

输入：

7 3
4 6 3 8 2 5 4

输出：

11

说明：

方案	Stage 0	Stage 1	Stage 2	最大负载
按层数尽量均分	[4,6,3]	[8,2]	[5,4]	$\max(13,10,9)=13$
最优	[4,6]	[3,8]	[2,5,4]	$\max(10,11,11)=11$

只有这一种划分能把峰值压到 $11$。

示例 2

输入：

4 2
2 8 3 5

输出：

10

说明：

四层耗时为 $[2,8,3,5]$，要切成 $2$ 个 Stage。

方案	Stage 0	Stage 1	负载	最大负载
A	[2]	[8,3,5]	$2$ vs $16$	$16$
B	[2,8]	[3,5]	$10$ vs $8$	$10$
C	[2,8,3]	[5]	$13$ vs $5$	$13$

最优为方案 B，峰值负载 $10$。

---

解题思路

第一步：理解问题本质

题目要求把 $N$ 层按原顺序切成 $K$ 段，每一段内的层编号连续、每一层恰好属于一段、 每段至少一层。一段的负载是段内各层前向时间之和，目标是让最大的那段负载尽量小。

这是典型的「最小化最大值」问题。它的三个要素是：

1. 切分方式爆炸：$N$ 层之间有 $N-1$ 个空隙，从中选 $K-1$ 个切开， 方案数是 $\binom{N-1}{K-1}$。$N=5000$ 时这个数字天文级别，不能枚举。
2. 答案有明确上下界：任何一层都不能被拆开，所以答案至少是 $\max f[i]$； 把所有层放在一段里，负载是总和，所以答案至多是 $\sum f[i]$。 答案一定落在这个区间里。
3. 「给定上限能否切出来」比「求最优值」容易得多：只要从左到右扫一遍， 每段尽量多装即可判断。于是可以把「求最优」转成「对答案二分」。

第二步：暴力解法

最直接的做法是枚举所有切点。用递归在 $N-1$ 个空隙上做「切 / 不切」的选择， 凑够 $K-1$ 个切点时计算一次最大段和：

best = float('inf')

def dfs(start, stages_left, current_max):
    """从 start 层开始，还剩 stages_left 段可用。"""
    global best
    if stages_left == 1:
        # 最后一段必须把剩下的层全部装下
        best = min(best, max(current_max, sum(f[start:])))
        return
    load = 0
    # 当前段的结束位置 end 从 start 开始，最多留 stages_left - 1 层给后面
    for end in range(start, n - stages_left + 1):
        load += f[end]
        dfs(end + 1, stages_left - 1, max(current_max, load))

这个做法能得到正确答案，但复杂度和方案数同阶。$N=5000$、$K=2500$ 时完全跑不完。

一个更进一步的过渡做法是动态规划：设 $dp[c][i]$ 表示「前 $i$ 层切成 $c$ 段」 的最小峰值，转移时枚举最后一段的起点：

$$dp[c][i]=\min_{0\le j<i}\ \max\left(dp[c-1][j],\ \text{sum}(f[j..i-1])\right)$$

它的时间复杂度是 $O(KN^2)$，$N=5000$ 时约 $6\times10^{10}$ 次运算，仍然太慢。 不过它揭示了一件重要的事：「给定上限能否切成 $K$ 段」的判定，比直接求最小值简单。

第三步：优化解法——把「求最优」变成「判可行」

先固定一个候选上限 $cap$，回答一个问题：

每一段的负载都不超过 $cap$ 时，最少要切多少段？

从左往右贪心：能装就装，装不下（再装一层就超过 $cap$）才切一刀。

segments = 1
current = 0
for f_i in f:
    if current + f_i <= cap:   # 还装得下，继续装
        current += f_i
    else:                      # 装不下，只能从这一层开始新的一段
        segments += 1
        current = f_i

这个贪心得到的是 $cap$ 下最少的段数。把任意一种合法切分与贪心逐段对比： 贪心的第一段一直装到装不下为止，所以它覆盖的前缀不会比任何合法方案的第一段短； 对第二段、第三段用同样的理由归纳下去，贪心在第 $t$ 段结束时覆盖的范围也不会落后。 既然每一段都尽量多装，它用的段数就是最少的。

于是判定规则是：

• 若 segments <= K，说明 $cap$ 可行。段数不够时，把任意一个含两层以上的段再切开即可， 因为 $K\le N$，总能切到正好 $K$ 段。
• 若 segments > K，说明 $cap$ 太小，无论怎么切都至少需要这么多段。

第四步：最优解法——二分答案

$cap$ 越大，贪心需要的段数越少（不会变多），所以「可行」这个性质关于 $cap$ 是单调的： 存在一个分界点 $ans$，$cap\ge ans$ 时可行，$cap<ans$ 时不可行。这正是二分的前提。

二分区间取 $[\max f[i],\ \sum f[i]]$：

• 下界取单层最大值，因为单层不可拆分，任何 $cap<\max f[i]$ 都不可能可行；
• 上界取总和，此时整条序列作为一段就满足要求，一定可行。

low, high = max(f), sum(f)
while low < high:
    mid = (low + high) // 2
    if segments_needed(f, mid) <= k:
        high = mid          # mid 可行，答案可能是 mid 或更小
    else:
        low = mid + 1       # mid 不可行，答案一定更大
return low                  # 区间收缩到一点，就是最小可行值

每一轮判定只扫一遍数组，二分的轮数是 $O(\log \sum f)$，总复杂度 $O(N\log\sum f)$。 $N=5000$、$\sum f\le 5\times10^6$ 时只有约 $5000\times23\approx 1.2\times10^5$ 次操作。

---

完整代码实现

import sys


def segments_needed(loads, cap):
    """在每段负载不超过 cap 的前提下，返回最少需要切成多少段。"""
    segments = 1          # 至少有一段，第一段从第 0 层开始
    current = 0           # 当前这一段的负载
    for f in loads:
        if current + f <= cap:
            # 还装得下就继续装：尽量少切，才能判断 cap 是否足够大
            current += f
        else:
            # 装不下了，必须从当前层开始新的一段
            segments += 1
            current = f
    return segments


def min_peak_load(loads, k):
    """二分最小的可行上限 cap，使得数组能切成不超过 k 段连续子数组。"""
    # 下界取单层最大值：任何一层都无法被拆开，所以 cap 不可能小于它。
    # 上界取总和：整段作为一段时负载就是总和，一定可行。
    low = max(loads)
    high = sum(loads)

    while low < high:
        mid = (low + high) // 2
        # 段数 <= k 说明 mid 够用，答案在左半边（含 mid）。
        # 段数 > k 说明 mid 太小，答案一定在 mid 右侧。
        if segments_needed(loads, mid) <= k:
            high = mid
        else:
            low = mid + 1

    return low


def solve():
    data = sys.stdin.buffer.read().split()
    if not data:
        return

    n = int(data[0])
    k = int(data[1])
    # 后面正好是 n 个层耗时；用切片一次取出，避免逐行读取的解析开销。
    loads = [int(x) for x in data[2:2 + n]]

    print(min_peak_load(loads, k))


if __name__ == "__main__":
    solve()

---

示例推演

样例 1

输入为 $N=7$、$K=3$，层耗时 $f=[4,6,3,8,2,5,4]$。

初始区间：$low=\max f=8$，$high=\sum f=32$。

轮次	$low$	$high$	$mid$	贪心切分	段数	是否可行	更新
1	8	32	20	[4,6,3] [8,2,5,4]	2	是	$high=20$
2	8	20	14	[4,6,3] [8,2] [5,4]	3	是	$high=14$
3	8	14	11	[4,6] [3,8] [2,5,4]	3	是	$high=11$
4	8	11	9	[4] [6,3] [8] [2,5] [4]	5	否	$low=10$
5	10	11	10	[4,6] [3] [8,2] [5,4]	4	否	$low=11$

第 5 轮结束后 $low=high=11$，输出 11。

第 4 轮的贪心过程可以对照着看：$cap=9$ 时，$4+6=10>9$ 所以要切； 接着 $6+3=9\le 9$ 可以装；$9+8=17>9$ 要切；$8+2=10>9$ 要切；$2+5=7\le 9$ 可以装； $7+4=11>9$ 要切。于是得到 5 段，超过 $K=3$，因此 $cap=9$ 不可行。

第 3 轮的 $cap=11$ 对应的正是题目说明里的最优方案 $[4,6]\ [3,8]\ [2,5,4]$，三段负载分别为 $10,11,11$，峰值为 $11$。

样例 2

输入为 $N=4$、$K=2$，层耗时 $f=[2,8,3,5]$。

初始区间：$low=8$，$high=18$。

轮次	$low$	$high$	$mid$	贪心切分	段数	是否可行	更新
1	8	18	13	[2,8,3] [5]	2	是	$high=13$
2	8	13	10	[2,8] [3,5]	2	是	$high=10$
3	8	10	9	[2] [8] [3,5]	3	否	$low=10$

$low=high=10$，输出 10，与题目说明里的方案 B 一致。

---

复杂度分析

解法	时间复杂度	空间复杂度	说明
枚举所有切点	$O\!\left(\binom{N-1}{K-1}\cdot N\right)$	$O(K)$	方案数爆炸，不可行
动态规划	$O(KN^2)$	$O(N)$	正确，但 $N=5000$ 时约 $6\times10^{10}$ 次运算
二分答案 + 贪心校验	$O(N\log\sum f)$	$O(N)$	本题采用，约 $1.2\times10^5$ 次操作

二分的轮数是 $O(\log(\sum f-\max f))$，每次判定只扫一遍数组，所以瓶颈就是这一遍扫描。 空间上除了保存 $N$ 个层耗时，只需要几个计数器。

---

易错点总结

1. 判定条件写成「段数恰好等于 $K$」

可行的条件是 segments <= K，不是 == K。段数比 $K$ 少时，只要某个段里有两层以上， 就可以继续切一刀把段数加一；因为 $K\le N$，总能切到正好 $K$ 段。 写成「等于」会把很多可行解误判成不可行。

2. 二分下界必须取单层最大值

如果 $cap$ 比某一层还小，那么这一层无论怎么放都会让所在段超载。 但贪心代码在这种情况下仍然会把它单独算成一段（段数看起来没有暴增）， 于是判定可能误判为可行。把下界设成 $\max f[i]$ 就从根上避免了这种输入。

3. 贪心不是「尽量平均分」

样例 1 已经说明：按层数平均分成 $[4,6,3]\ [8,2]\ [5,4]$ 的峰值是 $13$， 而正确答案是 $11$。贪心必须按负载判断「能不能再装一层」，而不是数层数。

4. 二分的区间写法

while low < high:
    mid = (low + high) // 2
    if feasible(mid):
        high = mid        # mid 可行，保留它
    else:
        low = mid + 1     # mid 不可行，跳过它

high = mid 而不是 high = mid - 1，因为 $mid$ 本身可能就是答案； low = mid + 1 而不是 low = mid，否则区间可能永远不收缩。

5. 边界情况

• $K=1$：答案是总和，二分会在 $high=\sum f$ 处收敛。
• $K=N$：答案是 $\max f[i]$，二分会在 $low=\max f$ 处收敛。
• $N=1$：同时满足上面两条，答案为这唯一一层的耗时。

6. 切分必须连续且非空

不能把层打乱重排（那不是流水线），也不能出现空 Stage。 贪心天然满足这两点：每次都从下一层开始新段，且只在「装不下」时才切。

---

扩展思考

• 如果每段还限制层数：例如每个 Stage 最多 $U$ 层，判定函数要同时数层数， 「装不下」的条件从「负载超限」变成「负载超限或层数超限」，二分框架不变。
• 如果层可以被拆开：那就从「连续切分」变成了「把工作量分给 $K$ 个工人」的调度问题， 判定变成「总工作量是否不超过 $K\cdot cap$」，贪心的形式完全不同。
• 如何输出具体方案：二分出最优峰值后，用同一个贪心再扫一遍记录切点即可。 若要求切点字典序最小（前面的段尽量长），贪心已经是这种形态。
• 和经典题的联系：本题与「分割数组的最大值」「复制书籍」是同一个模型， 都是「最小化最大连续段和」，标准解法都是二分答案 + 贪心校验。

---

相关题目

• 查找算法（二分答案：最小化最大值）
• 全局位与窗口注意力