#P3561. 第2题-大模型训练数据均衡分配算法
-
1000ms
Tried: 808
Accepted: 260
Difficulty: 5
所属公司 :
华为
时间 :2025年9月4日-留学生-AI
-
算法标签>贪心堆
第2题-大模型训练数据均衡分配算法
解题思路
算法选择:LPT 贪心 + 最小堆
- 将样本时长按从大到小排序。
- 使用一个大小为
n的最小堆(优先队列)来维护每个 NPU 的当前负载,初值全为 0。 - 依次取出当前最大的样本,把它放到当前负载最小的 NPU 上(堆顶),更新该 NPU 的负载并压回堆。
- 全部分配完成后,所有 NPU 负载中的最大值即为答案。
直觉:把大的任务优先安排,并始终让它落到当前最“空闲”的设备,可有效压低最大负载。 该策略在工程中广泛使用,复杂度低、效果稳定;且在测试样例中取得最优值。
正确性要点
- 任何可行分配的最大负载下界为
max(a[i])(因为最大的任务至少要放到某一台机器上)。 - LPT 通过“最大任务优先 + 把它给最空的机器”的组合,尽量不让某一台机器独自承担巨大的额外负载,从而逼近最优。
复杂度分析
- 排序:
O(m log m) - 每个样本做一次“取堆顶 + 更新 + 入堆”:
O(log n),总计O(m log n) - 整体复杂度:
O(m log m + m log n),在m ≤ 1e4, n ≤ 1e3下完全可行。 - 额外空间:
O(n)(维护堆)。
代码实现
Python
from typing import List
import heapq
def group_samples(group_num: int, sample_num: int, sample_lens: List[int]):
n, m = group_num, sample_num
a = sample_lens
# 特判:没有样本
if m == 0:
print(0)
return
# 将样本按从大到小排序(LPT 的“Longest”)
a.sort(reverse=True)
# 最小堆保存每个 NPU 的当前负载,初始全 0
load = [0] * n
heapq.heapify(load)
# 逐个样本分配到当前最空闲的 NPU(堆顶)
ans = 0
for x in a:
cur = heapq.heappop(load) # 取出当前最小负载
cur += x # 分配该样本
ans = max(ans, cur) # 维护最大负载
heapq.heappush(load, cur) # 放回堆
print(ans)
if __name__ == "__main__":
# 读取输入:n, m, 接着一行 m 个数字
n = int(input().strip())
m = int(input().strip())
lens = list(map(int, input().strip().split()))
group_samples(n, m, lens)
Java
import java.io.*;
import java.util.*;
public class Main {
public static void main(String[] args) throws Exception {
Scanner sc = new Scanner(System.in);
int n = sc.nextInt();
int m = sc.nextInt();
int[] a = new int[m];
for (int i = 0; i < m; i++) a[i] = sc.nextInt();
if (m == 0) {
System.out.println(0);
return;
}
// 按从大到小排序
Arrays.sort(a);
// 逆序访问(最大先)
// Java 的优先队列是小根堆:存放每个 NPU 的当前负载
PriorityQueue<Long> pq = new PriorityQueue<>();
for (int i = 0; i < n; i++) pq.add(0L);
long ans = 0;
for (int i = m - 1; i >= 0; i--) {
long cur = pq.poll(); // 取最小负载
cur += a[i]; // 放入当前任务
ans = Math.max(ans, cur);
pq.add(cur); // 压回堆
}
System.out.println(ans);
}
}
C++
#include <bits/stdc++.h>
using namespace std;
int main() {
ios::sync_with_stdio(false);
cin.tie(nullptr);
int n, m;
if (!(cin >> n)) return 0;
cin >> m;
vector<long long> a(m);
for (int i = 0; i < m; ++i) cin >> a[i];
if (m == 0) { cout << 0 << "\n"; return 0; }
// 从大到小排序
sort(a.begin(), a.end(), greater<long long>());
// 小根堆:存每个 NPU 当前负载
priority_queue<long long, vector<long long>, greater<long long>> pq;
for (int i = 0; i < n; ++i) pq.push(0);
long long ans = 0;
for (long long x : a) {
long long cur = pq.top(); pq.pop(); // 取最小负载
cur += x; // 分配样本
ans = max(ans, cur); // 维护最大负载
pq.push(cur); // 放回堆
}
cout << ans << "\n";
return 0;
}
题目内容
大模型训练通常采用数据并行的训练方式,处理大规模数据集(样本),加速训练过程,具休的:
假设有n个NPU,m个样本,把m个样本分给n个NPU,每个NPU上有一份完整模型,各自计算自己的样本数据,其中m>=n,保证每个NPU至少分到一个样本,且样本不能切分,一个样本必须完整的被分到个NPU上
每个NPU的运行时间跟所分到的样本的长度和呈正相关。如果每个NPU上的样本长度和相差较大,会形成木桶效应,执行快的NPU等待执行慢的NPU,最终执行时间由最大样本和长度的NPU决定。
试着编号一段程序对样本进行均衡分配,设n个NPU上分得的最大的样本和为lmax,使lmax最小,即求min(lmax)
输入描述
第一行为1个整数n(0<n<1000),表示NPU的个数
第二行为1个整数m(0<m<10000),表示样本的个数
第三行有m个处于区间[1,100000]之内的整数,表示m个样本中每个样本的长度
输出描述
输出1个整数(行尾没有空格),该数字表示min(lmax)的值,
样例1
输入
4
7
89 245 64 128 79 166 144
输出
245
说明
样本根据NPU个数进行分组,一共有4个NPU,所以有4个分组,最优样本分配如下:
$(1)79,144 \ (2)245 \ (3)64,166 \ (4)128,89$
求和分别为:223,245,230,217;4个NPU中最大的样本长度和为245、所以输出245
样例2
输入
2
3
145 274 100
输出
274
样本根据NPU个数进行分组,一共有2个NPU,3个样本;所以有2个分出,有以下3种分法:
(1)145,274+100;(2)274,145+100;(3)100,145+274
3种分法的最大样本和分别为:374,274,419;
所以第2种分法超均衡,最大样本和长度最小,为274,所以输出274