AGC023F - 01 on Tree

AGC023F - 01 on Tree

题目描述

Solution

有一个奇妙的贪心思路。（奇妙的原因是我不会证）
这一题的结点需要按拓扑序排序，并让逆序对个数最小。
考虑在儿子向父亲合并的过程中统计答案，产生的逆序对个数就是 $cnt[father][1]*cnt[son][0]$，其中 $cnt[x][0/1]$表示 $x$这个联通块已经拥有的 $0/1$结点的数量，我们贪心地选取 $\frac{cnt[x][0]}{cnt[x][1]}$最小的结点合并到父亲，并堆维护这一过程就可以做到 $O(nlgn)$求解答案。

#include <vector>
#include <list>
#include <map>
#include <set>
#include <deque>
#include <queue>
#include <stack>
#include <bitset>
#include <algorithm>
#include <functional>
#include <numeric>
#include <utility>
#include <sstream>
#include <iostream>
#include <iomanip>
#include <cstdio>
#include <cmath>
#include <cstdlib>
#include <cctype>
#include <string>
#include <cstring>
#include <ctime>
#include <cassert>
#include <string.h>
//#include <unordered_set>
//#include <unordered_map>
//#include <bits/stdc++.h>

#define MP(A,B) make_pair(A,B)
#define PB(A) push_back(A)
#define SIZE(A) ((int)A.size())
#define LEN(A) ((int)A.length())
#define FOR(i,a,b) for(int i=(a);i<(b);++i)
#define fi first
#define se second

using namespace std;

template<typename T>inline bool upmin(T &x,T y) { return y<x?x=y,1:0; }
template<typename T>inline bool upmax(T &x,T y) { return x<y?x=y,1:0; }

typedef long long ll;
typedef unsigned long long ull;
typedef long double lod;
typedef pair<int,int> PR;
typedef vector<int> VI;

const lod eps=1e-11;
const lod pi=acos(-1);
const int oo=1<<30;
const ll loo=1ll<<62;
const int mods=998244353;
const int MAXN=600005;
const int INF=0x3f3f3f3f;//1061109567
/*--------------------------------------------------------------------*/
inline int read()
{
	int f=1,x=0; char c=getchar();
	while (c<'0'||c>'9') { if (c=='-') f=-1; c=getchar(); }
	while (c>='0'&&c<='9') { x=(x<<3)+(x<<1)+(c^48); c=getchar(); }
	return x*f;
}
int cnt[MAXN][2],fa[MAXN],f[MAXN];
struct heapnode
{
	int x,y,id;
	bool operator < (const heapnode &a) const { return 1ll*x*a.y<1ll*y*a.x; }
};
priority_queue<heapnode> heap;
int find(int x){ return f[x]==x?f[x]:f[x]=find(f[x]); }
int main()
{
	int n=read();
	for (int i=2;i<=n;i++) fa[i]=read();
	for (int i=1;i<=n;i++) 
	{
		int x=read();
		cnt[i][x]++;
		f[i]=i;
	}
	for (int i=2;i<=n;i++) heap.push((heapnode){cnt[i][0],cnt[i][1],i});
	ll ans=0;
	while (!heap.empty())
	{
		int u=find(heap.top().id),x=heap.top().x,y=heap.top().y; heap.pop();
		if (cnt[u][0]==x&&cnt[u][1]==y)
		{
			int v=find(fa[u]);
			ans+=1ll*cnt[v][1]*cnt[u][0];
			cnt[v][0]+=cnt[u][0];
			cnt[v][1]+=cnt[u][1];
			f[u]=v;
			if (v!=1) heap.push((heapnode){cnt[v][0],cnt[v][1],v});
		}
	}
	printf("%lld\n",ans);
	return 0;
}