P5567. 激活检查点放置

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

题目描述

在大模型训练中，前向传播会留下大量中间激活，反向传播时还要用到它们。当 GPU 显存吃紧时，激活重计算（Gradient Checkpointing）是一种常见的「用时间换空间」做法：只在若干「检查点」层保存激活，其余层在反向阶段再算一遍。本题要求：在给定显存预算下，求出最优的 Checkpoint 放置方案。

模型共有 $N$ 层（编号为 $1$ 到 $N$），按前向顺序依次执行。给定两个长度为 $N$ 的正整数数组：

• forward_time[i]（$0 \le i < N$）：第 $i+1$ 层前向传播的耗时；
• memory[i]（$0 \le i < N$）：在第 $i+1$ 层保存激活（打 Checkpoint）所需的额外显存。

固定边界：

• 第 $0$ 层（输入层）的激活常驻显存，不计入额外开销，它是天然的起始 Checkpoint；
• 第 $N$ 层是终点：不必在第 $N$ 层再存额外激活，也不占用 Checkpoint 显存。

操作：你可以在某些中间层 $c$（$1 \le c \le N-1$）打 Checkpoint。在第 $c$ 层打 Checkpoint 会占用 memory[c-1] 的显存（第 $c$ 层对应数组下标 $c-1$）。所有 Checkpoint 的显存占用之和不得超过 $MaxMem$。

重计算代价：
设两个相邻 Checkpoint 位置为 $c$ 与 $d$（$c < d$，中间没有别的 Checkpoint）。反向传播经过该区间时，需要从第 $c$ 层的激活出发，重新前向计算以恢复中间层激活。由于反向是逐层推进的，层 $d-1$ 只需重算一次（从 $c$ 前向到 $d-1$），层 $d-2$ 需要重算两次，以此类推。因此该区间的重计算代价为：

$$Cost(c,d) = \sum_{k=c+1}^{d-1} forward\_time[k-1] \times (d-k)$$

其中 $forward\_time[k-1]$ 是第 $k$ 层的前向耗时，$(d-k)$ 是该层被重复计算的次数。若 $d = c+1$（相邻 Checkpoint），代价为 $0$。

目标：在显存预算 $MaxMem$ 内选择 Checkpoint 放置方案，使整网总重计算代价最小。总代价等于所有相邻 Checkpoint 对之间的代价之和（隐含起点 $0$ 与终点 $N$）。

输入格式

输入格式如下：

1. N MaxMem
2. f_1 f_2 ... f_N
3. m_1 m_2 ... m_N

• 第 $1$ 行：两个非负整数 $N$ 与 $MaxMem$，分别表示模型层数与额外显存预算上限。
• 第 $2$ 行：$N$ 个空格分隔的正整数，表示每层前向传播耗时。
• 第 $3$ 行：$N$ 个空格分隔的正整数，表示在每层保存 Checkpoint 所需的额外显存。

保证：$1 \le N \le 100$，$0 \le MaxMem \le 1000$，$1 \le f_i,m_i \le 1000$。

输出格式

输出一行一个整数：最小总重计算代价。

示例

示例 1

输入：

5 3
5 10 5 10 20
1 2 1 3 1

输出：

15

说明：

可选的 Checkpoint 位置：层 $1,2,3,4$。始终存在隐式起点 $0$ 与终点 $5$。

方案	Checkpoint 位置	显存消耗	区间划分	各区间代价	总代价
A	无	0	$(0,5)$	$5\cdot4+10\cdot3+5\cdot2+10\cdot1=70$	70
B	$\{2\}$	2	$(0,2),(2,5)$	$5$，$5\cdot2+10\cdot1=20$	25
C	$\{1,3\}$	2	$(0,1),(1,3),(3,5)$	$0$，$10$，$10$	20
D	$\{2,3\}$	3	$(0,2),(2,3),(3,5)$	$5$，$0$，$10$	15

方案 D 显存消耗 $3 \le 3$，总代价 $15$ 最小。

$Cost(0,5)$ 细算：

• $k=1$：$forward\_time[0]\times(5-1)=5\times4=20$
• $k=2$：$forward\_time[1]\times(5-2)=10\times3=30$
• $k=3$：$forward\_time[2]\times(5-3)=5\times2=10$
• $k=4$：$forward\_time[3]\times(5-4)=10\times1=10$
  合计 $70$。注意第 $5$ 层（$k=5$）不在求和范围内（上界为 $d-1=4$）。

提示

提示
可建模为带容量约束的区间最短路 / 分段 DP。
先计算 $cost[c][d]$（$0 \le c < d \le N$），利用前缀和可在 $O(N^2)$ 内完成：

• 先算 $S[i]=\sum_{k=1}^{i}forward\_time[k-1]$（前缀和）
• $Cost(c,d)=\sum_{k=c+1}^{d-1}forward\_time[k-1]\cdot(d-k)$，可用递推在 $O(1)$ 得到每对。

再构建转移：对每个位置 $i$ 与已用显存 $w$，枚举上一个 Checkpoint 位置 $j$（$0 \le j < i$）：

$$dp[i][w]=\min_{j}\bigl(dp[j][w-memory[i-1]]+cost[j][i]\bigr)$$

其中 $dp[0][0]=0$（起点不耗显存），第 $N$ 层作为终点不消耗额外显存。

---

解题思路

第一步：理解问题本质

把整条前向链路画成一条线：$0 \to 1 \to 2 \to \cdots \to N$。

• 第 $0$ 层是输入，激活本来就常驻显存，相当于免费的起点 Checkpoint；
• 第 $N$ 层之后不再需要保存激活，相当于免费的终点 Checkpoint；
• 中间的每一层 $c$ 都可以花 memory[c-1] 的显存，把它立成一个 Checkpoint。

于是一个方案就是沿着这条线挑几个切点。相邻两个切点 $c < d$ 之间形成一个段，反向传播时这一段要重算，代价为

$$Cost(c,d)=\sum_{k=c+1}^{d-1} forward\_time[k-1]\times(d-k)$$

为什么中间层会被算多次？反向传播是从第 $d$ 层的激活出发逐层往前推的：要恢复第 $d-1$ 层的激活，只需从 $c$ 往前向跑一次（算到 $d-1$）；但要恢复第 $d-2$ 层的激活，就得再跑一次前向（这次要跑到 $d-2$）；依此类推。于是第 $k$ 层被重复计算 $d-k$ 次，正好对应公式里的系数。

整题的目标是：

$$\min \sum_{\text{相邻切点 }(c,d)} Cost(c,d) \quad \text{s.t.} \quad \sum_{\text{选中的 } c} memory[c-1] \le MaxMem$$

这是典型的「区间代价 + 资源约束」结构：段的代价只取决于两个端点，彼此可以拼接。这类问题几乎总是可以化成「选切点」的 DP。

第二步：暴力解法

最直接的做法是枚举中间层的所有子集：

best = float("inf")
# 中间层一共有 N-1 个，用二进制枚举「打不打 Checkpoint」
for mask in range(1 << (n - 1)):
    marks = [0]                      # 起点总是切点
    used = 0
    for i in range(1, n):            # 第 i 层（1 起编号）
        if (mask >> (i - 1)) & 1:
            marks.append(i)
            used += memory[i - 1]
    marks.append(n)                  # 终点总是切点
    if used > max_mem:               # 显存超预算，方案非法
        continue

    total = 0
    for a, b in zip(marks, marks[1:]):
        for k in range(a + 1, b):    # 直接照抄题目公式
            total += forward_time[k - 1] * (b - k)
    best = min(best, total)

这个写法完全正确，但方案数是 $2^{N-1}$。当 $N=100$ 时是 $2^{99}$，远超任何评测机的时间限制，必须利用「段的代价可以复用」这一点做动态规划。

第三步：优化解法——先把所有区间代价算出来

DP 每做一次转移都要问「从 $c$ 直接跳到 $d$ 要花多少」。如果每次都现场求和，单次是 $O(N)$，整体会被拖到 $O(N^3\cdot MaxMem)$。先把所有 $Cost(c,d)$ 预处理成表，转移就变成 $O(1)$。

最朴素的做法是三重循环。更好的做法是递推：固定左端点 $c$，把右端点从 $d-1$ 拉长到 $d$。

对比两个式子：

$$Cost(c,d-1)=\sum_{k=c+1}^{d-2} forward\_time[k-1]\times(d-1-k)$$

$$Cost(c,d)=\sum_{k=c+1}^{d-1} forward\_time[k-1]\times(d-k)$$

当右端点从 $d-1$ 变成 $d$ 时：

1. 原来区间里的每一层 $k\in[c+1,d-2]$，系数从 $d-1-k$ 变成 $d-k$，都多加了一次，贡献 $\sum_{k=c+1}^{d-2}forward\_time[k-1]$；
2. 新进入区间的第 $d-1$ 层，系数是 $d-(d-1)=1$，贡献 $forward\_time[d-2]$。

两者合起来，正好就是 $forward\_time[c]+forward\_time[c+1]+\cdots+forward\_time[d-2]$，也就是

$$Cost(c,d)=Cost(c,d-1)+\bigl(prefix[d-1]-prefix[c]\bigr)$$

其中 $prefix[t]=forward\_time[0]+\cdots+forward\_time[t-1]$。注意下标对得很紧：第 $k$ 层的耗时是 forward_time[k-1]，所以「第 $c+1$ 层到第 $d-1$ 层」的耗时和是 $prefix[d-1]-prefix[c]$。

这样 $O(N^2)$ 就能填满整张表，而且 Cost(c, c+1) = 0（相邻 Checkpoint 之间没有中间层，不用重算）。

第四步：最优解法——以「最后一个 Checkpoint」为阶段的分段 DP

预处理完区间代价，剩下的问题就是：沿着 $0,1,\dots,N$ 挑切点，最小化各段代价之和，且切点的显存之和不超过 $MaxMem$。

定义状态：

$$dp[i][w]=\text{第 } i \text{ 层被打成 Checkpoint、且已占用 } w \text{ 显存时，}\ 0\to i \text{ 这一段的最小总代价}$$

这里的 $i$ 从 $0$ 开始编号，$i=0$ 就是免费的起点。

转移就是枚举上一个切点 $j$：

$$dp[i][w]=\min_{0\le j<i}\Bigl(dp[j]\bigl[w-memory[i-1]\bigr]+Cost(j,i)\Bigr)$$

要点有四个：

1. 起点单独初始化：$dp[0][0]=0$，其余全为 $+\infty$。起点不花显存。
2. 只有中间层才需要付显存：第 $i$ 层（$1\le i\le N-1$）打点是 memory[i-1]，所以 $i$ 只枚举到 $N-1$；终点 $N$ 不占显存，不进入 DP 的行。
3. $w$ 是「恰好占用」：$dp[i][w]$ 只对应「切点显存之和恰好等于 $w$」的方案。这样做不用在每个状态里再存「不超过 $w$ 的最优值」，最后在 $w\le MaxMem$ 上取最小即可。
4. 收尾单独处理：最后一段是从某个切点 $j$ 直达终点 $N$，代价是 $Cost(j,N)$，终点不占显存，所以

$$\text{answer}=\min_{0\le j<N,\ 0\le w\le MaxMem}\Bigl(dp[j][w]+Cost(j,N)\Bigr)$$

为什么不能贪心？因为「本段代价最小的切点」不一定让后面的段也便宜，而且显存是全局预算——在某一处多花一点显存换来的收益，取决于后面还有多长的段要切。所以必须让 DP 同时记住「切到哪」和「花了多少显存」两件事。

---

完整代码实现

"""
P5567. 激活检查点放置 - 区间代价 + 分段动态规划

输入格式：
- 第 1 行：两个整数 N、MaxMem；
- 第 2 行：N 个正整数 forward_time；
- 第 3 行：N 个正整数 memory。

输出格式：
- 一行一个整数：最小总重计算代价。

时间复杂度：O(N^2 * MaxMem)
空间复杂度：O(N * MaxMem + N^2)
"""

import sys


def build_costs(forward_time):
    """预处理任意两个 Checkpoint 之间的重计算代价 cost[c][d]。"""
    n = len(forward_time)

    # prefix[t] = forward_time[0] + ... + forward_time[t-1]
    # 第 k 层的耗时是 forward_time[k-1]，所以第 c+1..d-1 层的耗时和是
    # prefix[d-1] - prefix[c]。
    prefix = [0] * (n + 1)
    for t in range(1, n + 1):
        prefix[t] = prefix[t - 1] + forward_time[t - 1]

    cost = [[0] * (n + 1) for _ in range(n)]

    # 把右端点从 d-1 拉到 d：中间每一层的系数都加 1，
    # 相当于把 c+1..d-1 的耗时各再算一次，正好是一段前缀和。
    for c in range(n):
        for d in range(c + 2, n + 1):
            cost[c][d] = cost[c][d - 1] + prefix[d - 1] - prefix[c]

    return cost


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

    it = iter(data)
    n = int(next(it))
    max_mem = int(next(it))

    forward_time = [int(next(it)) for _ in range(n)]
    memory = [int(next(it)) for _ in range(n)]

    cost = build_costs(forward_time)

    INF = float("inf")

    # dp[i][w]：第 i 层作为最后一个 Checkpoint、已花费 w 显存时，
    # 从起点到第 i 层这一段的最小代价。
    dp = [[INF] * (max_mem + 1) for _ in range(n)]
    dp[0][0] = 0  # 起点免费

    for i in range(1, n):
        # 第 i 层打点的显存开销；它本身就超预算时，这一层不可能成为切点
        weight = memory[i - 1]
        if weight > max_mem:
            continue

        row = dp[i]
        for w in range(weight, max_mem + 1):
            best = INF
            # 枚举上一个切点 j，把 (j, i) 这一段的代价接上去
            for j in range(i):
                prev = dp[j][w - weight]
                if prev == INF:
                    continue
                candidate = prev + cost[j][i]
                if candidate < best:
                    best = candidate
            row[w] = best

    # 收尾：最后一段从切点 j 直达终点 n（终点不占显存）
    answer = INF
    for j in range(n):
        tail = cost[j][n]
        row = dp[j]
        for w in range(max_mem + 1):
            used = row[w]
            if used != INF and used + tail < answer:
                answer = used + tail

    print(answer)


if __name__ == "__main__":
    solve()

---

示例推演

输入：

5 3
5 10 5 10 20
1 2 1 3 1

即 $N=5$，$MaxMem=3$，forward_time = [5,10,5,10,20]，memory = [1,2,1,3,1]。

1）前缀和

$$prefix=[0,\ 5,\ 15,\ 20,\ 30,\ 50]$$

2）区间代价表

固定左端点，从左往右递推（只列出 $c<d$ 的格子）：

$c\backslash d$	1	2	3	4	5
0	0	5	20	40	70
1	—	0	10	25	50
2	—	—	0	5	20
3	—	—	—	0	10
4	—	—	—	—	0

以 $Cost(0,5)$ 为例，和题面完全一致：

$$Cost(0,4)=40,\qquad prefix[4]-prefix[0]=30-0=30,\qquad Cost(0,5)=40+30=70$$

再以 $Cost(2,5)$ 为例：$Cost(2,4)=forward\_time[2]\times(4-3)=5$，
加上 $prefix[4]-prefix[2]=30-15=15$，得到 $Cost(2,5)=5+15=20$。
按原始公式核对：$forward\_time[2]\times(5-3)+forward\_time[3]\times(5-4)=5\times2+10\times1=20$，两者一致。

3）DP 表（$w$ 为恰好占用的显存）

• 初始化：$dp[0][0]=0$，其余 $+\infty$。
• $i=1$（memory[0]=1）：只有 $dp[0][0]$ 可用，于是 $dp[1][1]=0+Cost(0,1)=0$。
• $i=2$（memory[1]=2）：
  • $w=2$：$dp[0][0]+Cost(0,2)=0+5=5$；
  • $w=3$：$dp[1][1]+Cost(1,2)=0+0=0$。
• $i=3$（memory[2]=1）：
  • $w=1$：$dp[0][0]+Cost(0,3)=20$；
  • $w=2$：$dp[1][1]+Cost(1,3)=0+10=10$；
  • $w=3$：$dp[2][2]+Cost(2,3)=5+0=5$。
• $i=4$（memory[3]=3）：只有 $w=3$ 可行，$dp[0][0]+Cost(0,4)=40$。

$i\backslash w$	0	1	2	3
0	0	∞	∞	∞
1	∞	0	∞	∞
2	∞	∞	5	0
3	∞	20	10	5
4	∞	∞	∞	40

4）收尾取最小值

对每个 $j$ 与 $w$ 计算 $dp[j][w]+Cost(j,5)$：

$j$	可行的 $w$	$dp[j][w]$	$Cost(j,5)$	合计
0	0	0	70	70
1	1	0	50	50
2	2 / 3	5 / 0	20	25 / 20
3	1 / 2 / 3	20 / 10 / 5	10	30 / 20 / 15
4	3	40	0	40

最小是 $15$，来自 $j=3,\ w=3$。回溯它的来源：$dp[3][3]=dp[2][2]+Cost(2,3)=5+0$，而 $dp[2][2]=dp[0][0]+Cost(0,2)=5$，所以方案是「$0\to 2\to 3\to 5$」，即第 $2$、$3$ 层打 Checkpoint：

• 显存：$memory[1]+memory[2]=2+1=3\le MaxMem=3$；
• 代价：$Cost(0,2)+Cost(2,3)+Cost(3,5)=5+0+10=15$。

与题面表格里的方案 D 完全一致，输出 15。

---

复杂度分析

解法	时间复杂度	空间复杂度	说明
枚举所有 Checkpoint 子集	$O(2^{N-1}\cdot N)$	$O(N)$	$N=100$ 时不可行
区间代价（三重循环）+ 分段 DP	$O(N^3+N^2\cdot MaxMem)$	$O(N\cdot MaxMem+N^2)$	正确，但区间代价可以更快算
区间代价（递推 + 前缀和）+ 分段 DP	$O(N^2\cdot MaxMem)$	$O(N\cdot MaxMem+N^2)$	本题采用

本题 $N\le100$、$MaxMem\le1000$，转移量约 $\frac{N^2}{2}\cdot MaxMem\approx5\times10^6$，在 Python 中也能秒级完成。

---

易错点总结

1. 求和范围是 $c+1$ 到 $d-1$

区间的两个端点是 Checkpoint，它们的激活是存下来的，不需要重算；真正被重算的是夹在中间的 $c+1,\dots,d-1$。写成 range(c, d) 会把 $c$ 自己也按 $d-c$ 次算进去，代价必然偏大。这类下标错误不会报异常，只会让答案悄悄变大，所以一定要盯着样例核对一遍。

2. 下标错位一格

题目用 $1$ 起编号说层，数组是 $0$ 起下标：第 $k$ 层的耗时是 forward_time[k-1]，第 $c$ 层的显存是 memory[c-1]。DP 里 $i$ 是「第 $i$ 层」，所以扣的是 memory[i-1]。

3. 终点不能被选成 Checkpoint

第 $N$ 层是免费终点，不花显存。所以 DP 的行只到 $N-1$，最后一段单独用 $Cost(j,N)$ 收尾。如果把 $N$ 也当成普通层去扣 memory[N-1]，会白白多花显存甚至判成无解。

4. 起点要单独初始化

$dp[0][0]=0$，$dp[0][w>0]=+\infty$。起点是「天然存在的 Checkpoint」，既不花显存也不需要重算代价。

5. 显存是「不超过」而不是「恰好等于」

$dp[i][w]$ 记录的是恰好用掉 $w$ 的方案，但答案要在所有 $w\le MaxMem$ 上取最小。只读 $dp[j][MaxMem]$ 会漏掉「用了更少显存、代价却更低」的方案：显存预算 $10$、最优解只花了 $1$ 时，dp[j][10] 可能还是 $+\infty$，直接输出就会出错。

6. 不可达状态必须是 $+\infty$

if weight > max_mem: continue 与 if prev == INF: continue 两处都不能省：前者保证超预算的层不会成为切点，后者保证不会把「不存在的方案」当成代价 $0$ 参与比较。

7. $N=1$ 的边界

只有输入层和终点，没有中间层可选，起点直接连到终点，$Cost(0,1)=0$，答案就是 0。

8. 全部层都能打 Checkpoint 时答案是 0

相邻 Checkpoint 之间没有中间层，$Cost(c,c+1)=0$。此时 DP 会自然给出 $0$，不需要额外特判。

---

扩展思考

• 要求输出方案本身：再开一张 pre[i][w] 记录转移来的 $j$，从最优终点倒推即可还原所有 Checkpoint 位置。
• 改成「最多打 $k$ 个 Checkpoint」：把状态加一维 dp[i][cnt][w]，或者干脆不要显存维度、只按个数分层，复杂度 $O(kN^2)$。这类「切 $k$ 刀」的题（如「山区建小学」）与本题是同一族。
• 显存上限很大时：$O(N^2\cdot MaxMem)$ 会变慢。这时常见做法是拉格朗日松弛（给显存乘一个价格，把约束折进目标函数）或直接做 Pareto 前沿，但要注意这些做法只能给出近似或需要额外验证。
• 时间/空间的权衡直觉：Checkpoint 越密，重计算越少、显存越多；越稀则相反。本题其实就是在这条曲线上找「预算内最省时间」的那个点。
• 区间代价的另一种理解：$Cost(c,d)$ 也可以写成 $\sum_{t=1}^{d-c-1}\sum_{k=c+1}^{d-t}forward\_time[k-1]$，即「每一轮重算所覆盖的层」之和，与递推式的推导是同一件事。

---

相关题目

• NPU梯度任务均衡调度
• 双资源约束下的运维任务收益最大化
• 动态规划