FII Code 2021 Round #2 : E. Scrambled Eggs

なんだこのタイトル。
https://csacademy.com/contest/fii-code-2021-round-2/task/scrambled-eggs/

問題

正整数がN個与えられる。
このうちK個を選び、それぞれ素因数を1個ずつ選ぼう。
この時、各素因数の選出回数の最大値がXとなるようにしたい。
選び方の一例を示せ。

解法

まず入力値を素因数分解しよう。
その際、X回以上登場する素因数pがあれば、そのうち1つを「必ずX回登場すべき素因数」に指定する。
証明はEditorialに任せるとして、pはX回以上登場するならどれを選択してもよい。
他の素因数は、「X回以下登場すべき素因数で、pを除く全素因数で合計K-X回登場すべき」とする。

あとは、上記内容をグラフの最大フロー問題に落とし込み、復元しよう。
これは以下のようなグラフを作る。

• sourceから、各入力値に対応する点に容量1の辺を張る。
• 各素因数に対応する点を作り、各入力値のうちその素因数を持つものについて、対応する点の間に容量1の辺を張る
• pに対応する点からsinkに容量Xの辺を張る
• p以外の素因数に対応する点から、点Vに容量Xの辺を張る。V→sinkに容量(K-X)の辺を張る

int N,K,X;
int A[101010];
vector<int> V[101010];
const int prime_max = 1000000;
vector<int> prime;
int NP,divp[prime_max];

void cprime() {
	if(NP) return;
	for(int i=2;i<prime_max;i++) if(divp[i]==0) {
		//M[i]=NP;
		prime.push_back(i); NP++;
		for(ll j=1LL*i*i;j>=i&&j<prime_max;j+=i) if(divp[j]==0) divp[j]=i;
	}
}

template<class V> class MaxFlow_dinic {
public:
	struct edge { int to,reve;V cap;};
	static const int MV = 2020202;
	vector<edge> E[MV];
	int itr[MV],lev[MV];
	void add_edge(int x,int y,V cap,bool undir=false) {
		E[x].push_back((edge){y,(int)E[y].size(),cap});
		E[y].push_back((edge){x,(int)E[x].size()-1,undir?cap:0});
	}
	void bfs(int cur) {
		MINUS(lev);
		queue<int> q;
		lev[cur]=0;
		q.push(cur);
		while(q.size()) {
			int v=q.front(); q.pop();
			FORR(e,E[v]) if(e.cap>0 && lev[e.to]<0) lev[e.to]=lev[v]+1, q.push(e.to);
		}
	}
	V dfs(int from,int to,V cf) {
		if(from==to) return cf;
		for(;itr[from]<E[from].size();itr[from]++) {
			edge* e=&E[from][itr[from]];
			if(e->cap>0 && lev[from]<lev[e->to]) {
				V f=dfs(e->to,to,min(cf,e->cap));
				if(f>0) {
					e->cap-=f;
					E[e->to][e->reve].cap += f;
					return f;
				}
			}
		}
		return 0;
	}
	V maxflow(int from, int to) {
		V fl=0,tf;
		while(1) {
			bfs(from);
			if(lev[to]<0) return fl;
			ZERO(itr);
			while((tf=dfs(from,to,numeric_limits<V>::max()))>0) fl+=tf;
		}
	}
};

map<int,int> P,Q;
MaxFlow_dinic<int> mf;
void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cprime();
	
	cin>>N>>K>>X;
	FOR(i,N) {
		cin>>A[i];
		for(j=0;prime[j]*prime[j]<=A[i];j++) if(A[i]%prime[j]==0) {
			V[i].push_back(prime[j]);
			P[prime[j]]++;
			while(A[i]%prime[j]==0) A[i]/=prime[j];
		}
		if(A[i]>1) {
			V[i].push_back(A[i]);
			P[A[i]]++;
		}
	}
	x=0;
	y=-1;
	vector<int> QQ;
	FORR2(p,q,P) {
		QQ.push_back(p);
		Q[p]=x++;
		if(q>=X) y=p;
	}
	if(y==-1) {
		cout<<"IMPOSSIBLE"<<endl;
		return;
	}
	FOR(i,N) {
		mf.add_edge(0,i+1,1);
		FORR(v,V[i]) mf.add_edge(i+1,(N+1)+Q[v],1);
	}
	FORR2(p,q,Q) {
		if(p==y) {
			mf.add_edge((N+1)+q,1500000,X);
		}
		else {
			mf.add_edge((N+1)+q,1500001,X);
		}
	}
	mf.add_edge(1500001,1500000,K-X);
	x=mf.maxflow(0,1500000);
	
	if(x<K) {
		cout<<"IMPOSSIBLE"<<endl;
	}
	else {
		FOR(i,N) {
			int tar=-1;
			FORR(e,mf.E[i+1]) if(e.to>0&&e.cap==0) {
				
				tar=QQ[e.to-(N+1)];
			}
			cout<<tar<<" ";
		}
		cout<<endl;
		
	}
	
}

まとめ

pがなんでも良いって点に気付けるかが重要か。
あと頂点数・辺数が多そうだけど、これで間に合うのね。