【AT Keyence2019E】Connecting Cities

题目

题目链接：https://atcoder.jp/contests/keyence2019/tasks/keyence2019_e
有 \(n\) 个点排成一行，在第 \(i,j\) 个点之间连边的代价为 \(|i-j| \times D+A_i+A_j\)，求将它们连成一棵树的最小代价。
\(n\leq 2\times 10^5\)。

思路

麻了这种题真的是给人想的吗。
直接建边跑最小生成树的边数是 \(O(n^2)\) 的，考虑怎样才能减少边数。
对于一个区间 \([l,r]\)，考虑跨过中点 \(mid\) 的边，区间 \([l,mid]\) 中点 \(i\) 的代价为 \(val_i=a_i-d\times i\)，区间 \((mid,r]\) 中点 \(j\) 的代价为 \(val_j=a_j+d\times j\)。
对于 \(i,i'\in [l,mid]\) 且 \(val_i<val_{i'}\)，以及 \(j,j'\in(mid,r]\) 且 \(val_j<val_j'\)，则 MST 中不可能有 \((i',j')\) 这条边。因为 MST 中任意一条边不可能是原图中某一个环的最大边（否则可以断掉这条边，连上环的另一条边更优），而在 \((i,j),(i',j'),(i,j'),(i',j)\) 这四条边组成的环中，\((i',j')\) 的权值是最大的。
这样的话连接左右两个区间的边中，至少有一个端点是那一边的最小值。那么一个长度为 \(len\) 的区间只有 \(O(len)\) 条边可选。继续分治下去，这张图就化简成只有 \(O(n\log n)\) 条边了。
时间复杂度 \(O(n\log ^2 n)\)。

代码

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;

const int N=200010,LG=18,Inf=1e9;
int n,m,father[N];
ll ans,d,a[N];

struct edge
{
	int u,v;
	ll dis;
}e[N*LG];

void solve(int l,int r)
{
	if (l==r) return;
	int mid=(l+r)>>1,p=l;
	for (int i=l;i<=mid;i++)
		if (a[i]-d*i<=a[p]-d*p) p=i;
	for (int i=mid+1;i<=r;i++)
		e[++m]=(edge){p,i,a[p]+a[i]+d*(i-p)};
	p=r;
	for (int i=mid+1;i<=r;i++)
		if (a[i]+d*i<=a[p]+d*p) p=i;
	for (int i=l;i<=mid;i++)
		e[++m]=(edge){p,i,a[p]+a[i]+d*(p-i)};
	solve(l,mid); solve(mid+1,r);
}

bool cmp(edge x,edge y)
{
	return x.dis<y.dis;
}

int find(int x)
{
	return x==father[x]?x:father[x]=find(father[x]);
}

void kruskal()
{
	for (int i=1;i<=n;i++) father[i]=i;
	sort(e+1,e+1+m,cmp);
	for (int i=1;i<=m;i++)
	{
		int x=find(e[i].u),y=find(e[i].v);
		if (x!=y) father[x]=y,ans+=e[i].dis;
	}
}

int main()
{
	scanf("%d",&n);
	scanf("%lld",&d);
	for (int i=1;i<=n;i++) scanf("%lld",&a[i]);
	solve(1,n);
	kruskal();
	cout<<ans;
	return 0;
}