华为笔试华为AI算法 0903

笔试时间：2025年9月3日

往年笔试合集：

2023春招秋招笔试合集

2024春招秋招笔试合集

第一题

在 MOE 模型训练时，token 会依据概率发送到 topk 个不同专家计算，这些专家分布在多个 NPU 卡上。Device - Limited routing 算法可将 token 路由目标限制在 p个 NPU 以降低通信成本，具体步骤如下：

把 n个专家平均分配在 m个 NPU 上，每个 NPU 上的专家为一个组；设 n个专家的编号为 N = [0,1,2,…,n−1]，同一个专家组上的专家编号是连续的；

每个专家对应一个概率，表示被路由到的可能性；用每个组中的最大概率作为本组代表，从所有组中选择概率最大的 p个组，其所在的 NPU 即为路由目标限制 NPU；

再从上述 p个 NPU 对应的所有专家概率中选择 k个最大的概率对应的专家编号作为最终路由目标。

试编写一段程序，实现以上路由算法。

输入描述

第一行有 4 个处于区间 [1,10000]之内的整数，第 1 个表示专家的个数 n，第 2 个表示 NPU 个数 m，第 3 个表示路由目标限制 NPU 数 p，第 4 个表示目标路由专家个数 k；

第二行有 n个处于区间 (0,1)之内的浮点数，表示每个专家对应的概率值，这 n个数对应的专家的编号为 [0,1,2,…,n−1]。

输出描述

如果 n不能被 m整除或者获取不到 k个专家编号，输出 error；

否则，按照从小到大的顺序，输出 k个专家编号，任意相邻两数之间有空格，最后一个数字（行尾没有空格）。

样例输入

8 4 4 2 0.5 0.01 0.09 0.023 0.027 0.05 0.1 0.2

样例输出

0.7

参考题解

模拟 Device-Limited routing的流程。

Python：

def solve():
    n, m, p, k = map(int, input().split())
    a = list(map(float, input().split()))

    # 校验
    if 0 != n % m:
        print("error"); return
    if m < p:
        print("error"); return
    q = n // m
    if k > p * q:
        print("error"); return

    # 分组（改为 while）
    g, i = [], 0
    while i < m:
        s = i * q
        e = s + q
        t, j = [], s
        while j < e:
            t.append((a[j], j))
            j += 1
        g.append(t)
        i += 1

    # 选 p 个组（按组内最大概率，等价于原逻辑；用二元键确保与稳定排序一致）
    r, i = [], 0
    while i < m:
        mx = max(g[i])[0]
        r.append((mx, i))
        i += 1
    r.sort(key=lambda z: (-z[0], z[1]))  # 概率降序，组编号升序用于平手打破

    sel = [idx for _, idx in r[:p]]

    # 汇总候选并取前 k（Python 排序稳定，按概率降序与原行为等价）
    pool = []
    for idx in sel:
        pool.extend(g[idx])
    pool.sort(key=lambda x: x[0], reverse=True)

    ans = [j for _, j in pool[:k]]
    ans.sort()
    print(" ".join(map(str, ans)))


if __name__ == "__main__":
    solve()