题目内容

在混合专家模型（Mixture of Experts, MoE）中，传统 Token Choice 策略由每个 token 选择专家，容易造成热门专家过载、冷门专家闲置的负载不均衡问题。Expert Choice Routing（ECR）反转选择关系：每个专家主动选择容量范围内的 token，使各专家的处理量天然保持均衡。

本题要求在不调用 pytorch、tensorflow 等深度学习库的前提下，实现 ECR 机制。给定 token 表示矩阵 $X$、门控权重矩阵 $W_g$ 以及每个专家的变换矩阵 $W_e$，需要完成以下计算。

对每个 token $i$ 和每个专家 $e$，计算门控分数：

再对同一 token 在所有专家上的分数做数值稳定 softmax，得到归一化门控权重：

之后，每个专家 $e$ 独立对所有 token 排序。排序键为先按 $g_{i,e}$ 从大到小，若 $g_{i,e}$ 相同则按 token 编号 $i$ 从小到大。每个专家只保留排序后的前 $S$ 个 token。

对被专家 $e$ 选中的 token $i$，计算其行向量 $X_i$ 与专家矩阵 $W_e$ 的乘积，再乘以 $g_{i,e}$，并将结果累加到该 token 的最终表示中。若某个 token 没有被任何专家选中，则其最终表示保持为长度 $d$ 的全 0 向量。

这里 $N$ 是 token 数量，$d$ 是特征维度，$E$ 是专家数量，$S$ 是每个专家的 token 容量且不超过 $N$。

输入描述

程序从标准输入读取数据，所有数值以空白字符分隔。

第一行包含 4 个整数 $N$、$d$、$E$、$S$，分别表示 token 数量、特征维度、专家数量、每个专家选择的 token 容量。

接下来是输入 token 表示矩阵 $X$，共 $N$ 行，每行 $d$ 个浮点数。

然后是门控权重矩阵 $W_g$，共 $d$ 行，每行 $E$ 个浮点数。该矩阵按转置形式给出：行下标表示输入特征维度，列下标表示专家编号。

最后依次是 $E$ 个专家块。每个专家块包含 $d$ 行，每行 $d$ 个浮点数，表示一个 $d\times d$ 的专家变换矩阵 $W_e$。

输出描述

若标准输入中的数值个数不足以读入首行的 4 个整数，或者首行参数不满足 $N>0$、$d>0$、$E>0$、$S>0$、$S\le N$，输出单个整数 0。

否则，输出 $N$ 行，每行包含 $d$ 个浮点数，表示最终表示矩阵 $Y$。每个浮点数保留 2 位小数，同一行内用单个空格分隔。

样例1

输入

0 2 2 1

输出

0

说明

输入的首行参数中 $N$ 的值为 0，不满足 $N>0$ 的约束。因此程序输出整数 0。

样例2

输入

1 1 1 1
2.0
0.5
3.0

输出

6.00

说明

参数合法：$N=1$，$d=1$，$E=1$，$S=1$。只有 1 个 token。

门控分数为 $2.0\times0.5=1.0$。因为只有一个专家，softmax 归一化后的权重为 $1.0$。专家容量为 1，选中该 token。

最终结果为 $1.0\times2.0\times3.0=6.0$，保留两位小数输出 6.00。

样例3

输入

3 1 2 1
0.0
1.0
0.5
0.0 1.0
1.0
1.0

输出

0.00
0.73
0.00

说明

对 3 个 token 分别计算门控分数和 softmax 权重。token 0 的分数为 $(0,0)$，所以 $g_{0,0}=g_{0,1}=0.5$。token 1 的分数为 $(0,1)$，所以 $g_{1,0}\approx0.2689$、$g_{1,1}\approx0.7311$。token 2 的分数为 $(0,0.5)$，所以 $g_{2,0}\approx0.3775$、$g_{2,1}\approx0.6225$。

专家 0 按 $g_{i,0}$ 选择最大 token，最大值为 token 0 的 0.5，因此专家 0 选择 token 0。专家 1 按 $g_{i,1}$ 选择最大 token，最大值为 token 1 的 0.7311，因此专家 1 选择 token 1。

token 0 被专家 0 选中，贡献为 $0.5\times0.0\times1.0=0.0$，输出 0.00。token 1 被专家 1 选中，贡献为 $0.7311\times1.0\times1.0\approx0.7311$，保留两位为 0.73。token 2 未被任何专家选中，输出保持 0.00。