AtCoder Beginner Contest 257 Ex

这么难的题出在ABC真的好吗？

首先是一波推式子，假设我们选的骰子集合为$D=\{i1,i2,...,ik\}$（这里$D$储存的是骰子的下标），那么我们不难列出下面的式子：

$$E_{D}=\frac{1}{6^k}\sum_{j1,j2,...,jk} (A_{i1,j1}+A_{i2,j2}+..+A_{ik,jk})^2-(C_{i1}+C_{i2}+...+C_{ik})$$

考虑前半部分很烦，我们可以拆掉平方，即利用完全平方公式

$$(a_{1}+a_{2}+...+a_{{n}})^2=\sum_{i=1}^{n} a_{i}^2+2\sum_{i=1}^{n-1}\sum_{j=i+1}^{n} a_{i}a_{j}$$

可以得到：

$$E_{D}=\frac{1}{6^k}\sum_{j1,j2,...,jk}( \sum {A_{i,j}}^2+\sum_{i\in D,1\leq j\leq 6,k\in D,i<k,1\leq l\leq 6} 2A_{i,j}A_{k,l})-(C_{i1}+C_{i2}+...+C_{ik})$$

我们发现，对于$i\in D,1\leq j\leq 6$，$A_{i,j}^2$的贡献是$6^{k-1}$（相当于固定了这一位，其它$(k-1)$个骰子可以在$1$到$6$中任意选择），而$A_{i,j}A_{k,l}$的贡献是$6^{k-2}$（相当于固定这$2$位，其它$(k-1)$个骰子任意选择）

所以，我们就可以愉快的把最外面，也是最难枚举的$\sum_{j1,j2,...,jk}$给去除了！

$$E_{D}=\frac{1}{6^k} \sum_{i\in D} 6^{k-1} {A_{i,j}}^2 +\frac{1}{6^k} \sum_{i,k\in D,i<k} 2\times 6^{k-2}A_{i,j}A_{k,l} -\sum_{i\in D} C_{i}$$

将$\frac{1}{6^k}$乘进去：

$$E_{D}=\frac{1}{6}\sum_{i\in D}{A_{i,j}}^2 +\frac{1}{18} \sum_{i,k\in D,i<k} A_{i,j}A_{k,l} -\sum_{i\in D} C_{i}$$

我们发现：

$$\sum_{i\in D}\sum_{k\in D,i<k} A_{i,j}A_{k,l}=\frac{(\sum_{i\in D} \sum_{j=1}^{6}A_{i,j})(\sum_{i\in D} \sum_{j=1}^{6} A_{i,j})-\sum_{i\in D} \sum_{j=1}^{6}{A_{i,j}}^2}{2}$$

（相当于两两组合，但是每一对$(i,j)$要被计算两遍：在$(i,j)$和$(j,i)$时都会被计算，故要除以$2$，并且还要减去自己和自己组合的情况）

那么我们设$S_{i}=\sum_{j=1}^{6} A_{i,j}$，$P_{i}=\sum_{j=1}^{6} A_{i,j}^2$，可得：

$$E_{D}=\frac{1}{6}\sum_{i\in D} P_{i} + \frac{1}{36} (\sum_{i\in D} S_{i})^2-\frac{1}{36} \sum_{i\in D} {S_{i}}^2-\sum_{i\in D} C_{i}$$

设$T_{i}=6P_{i} - {S_{i}}^2 - 36C_{i}$。

则：

$$E_{D}=\frac{1}{36}(\sum_{i\in D} S_{i})^2 + \frac{1}{36} \sum_{i\in D} T_{i}$$

$S_{i},T_{i}$可以在输入时就计算好。那么现在我们就将原问题转成如下的问题：

给你$n$个数对$(x_{i},y_{i})$，让你选出$k$个数对，使得$(\sum x_{i})^2+\sum y_{i}$最大。

此时我们可以得到一个非常重要的结论：

若我们选择集合$D$最优，那么一定存在一个整数$c$，满足对于任意$i\in D,j\notin D$，$cx_{i}+y_{i}\geq cx_{j}+y_{j}$。

证明：

咕咕咕（马上补）

那么此时我们就可以枚举整数$c$从$-\inf$到$\inf$，找到当前$cx_{i}+y_{i}$最大的$k$个，将$(\sum x)^2+\sum y$更新到答案。

此时时间复杂度为$O(n\log n\times \inf)$，虽然不能AC，但是已经不是一个指数级别的做法了。我们已经有了很大的进步。

其实还可以优化。我们利用类似离散化的思想：如果$c$时刻的$cx_{i}+y_{i}$的数列和$(c+1)$时刻的不变，那么我们没有必要枚举$(c+1)$时刻。所以，数列只会在从某对$(i,j)$从$cx_{i}+y_{i} < cx_{j}+y_{j}$变成$cx_{i}+y_{i} \geq cx_{j}+y_{j}$时改变。即$c=\lceil \frac{y_{j}-y_{i}}{x_{i}-x_{j}}\rceil$时改变。那么我们把这$O(n^2)$个时刻存储下来，枚举$c$等于它们的时候，以及最初（$c=-\inf$）时即可。时间复杂度为$O(n^3\log n)$。

继续优化！我们发现上面做法慢的原因在于每次枚举新的$c$时，总是要对数组重新排序。如果我们可以利用单调性维护排序数组呢？我们发现，当$c=\lceil \frac{y_{j}-y_{i}}{x_{i}-x_{j}}\rceil$时，我们只改变$i,j$的大小关系，也就是说对于之前的数组，我们只要交换$i,j$就可以得到现在的数组。故我们就可以先排序，之后每枚举一个$c$，交换对应的下标即可，时间复杂度为$O(n^2\log n)$。

代码如下：

#include<bits/stdc++.h>
#define debug(...) std::cerr<<#__VA_ARGS__<<" : "<<__VA_ARGS__<<std::endl

using ll=long long;

const int maxn=1005;
const ll inf=4e18;
int n,k;
int c[maxn],a[maxn][10],id[maxn];
ll s[maxn],t[maxn],p[maxn];

ll calc(int ind) {
	int x=id[ind],y=id[ind+1];
	if(s[x]>=s[y]) return inf;
	return (t[x]-t[y]+s[y]-s[x]-1ll)/(s[y]-s[x]);
}

int main() {
	scanf("%d%d",&n,&k);
	for(int i=1;i<=n;i++) scanf("%d",&c[i]);
	for(int i=1;i<=n;i++)
		for(int j=1;j<=6;j++)
			scanf("%d",&a[i][j]);
	ll Ss=0,St=0,ans=-inf;
	for(int i=1;i<=n;i++) {
		for(int j=1;j<=6;j++) {
			s[i]+=1ll*a[i][j];
			p[i]+=1ll*a[i][j]*a[i][j];
		}
		t[i]=6ll*p[i]-s[i]*s[i]-36ll*c[i];
		id[i]=i;
	}
	std::sort(id+1,id+n+1,[](int id1,int id2){return t[id1]>t[id2];});
	for(int i=1;i<=k;i++) Ss+=s[id[i]],St+=t[id[i]];
	ans=std::max(ans,Ss*Ss+St);
	std::set<std::pair<long long,int>> S;
	for(int i=1;i<n;i++) S.insert({calc(i),i});
	while(!S.empty()) {
		std::pair<long long,int> rec=*S.begin();
		if(rec.first==inf) break;
		S.erase({calc(rec.second),rec.second});
		if(rec.second>=2) {
			S.erase({calc(rec.second-1),rec.second-1});
		}
		if(rec.second<=n-2) {
			S.erase({calc(rec.second+1),rec.second+1});
		}
		if(rec.second==k) {
			Ss-=s[id[rec.second]]; St-=t[id[rec.second]];
		}
		std::swap(id[rec.second],id[rec.second+1]);
		S.insert({calc(rec.second),rec.second});
		if(rec.second>=2) {
			S.insert({calc(rec.second-1),rec.second-1});
		}
		if(rec.second<=n-2) {
			S.insert({calc(rec.second+1),rec.second+1});
		}
		if(rec.second==k) {
			Ss+=s[id[rec.second]]; St+=t[id[rec.second]];
		}
		ans=std::max(ans,Ss*Ss+St);
	}
	//注意ans还有可能为负数
	ans=(ans%998244353+998244353)%998244353;
	printf("%lld\n",ans*859599304ll%998244353ll);
	return 0;
}