kmjp's blog

競技プログラミング参加記です

AtCoder ARC #138 (大和証券プログラミングコンテスト2022 Spring) : F - KD Tree

1つ1つの考察ステップは難しくないんだけどもね…。
https://atcoder.jp/contests/arc138/tasks/arc138_f

問題

1~Nの順列Pが与えられる。
数値のペアの配列[(1,P[1]),(2,P[2]).....(N,P[N])]を、以下の手順で並べ替える。

  • 1つ目の要素または2つ目の要素に対し値を指定する。そして、数列の各要素を、指定した値未満のものと以上のものの2つの数列に分解する。
  • 両数列それぞれについて再帰的にこの操作を行ったうち、得られた結果を「指定した値未満のもの」「以上のもの」の順で連結する

この手順で得られる並べ替え後の数列は何通りか。

解法

メモ化再帰で解く。
基本的なアプローチは、分解のパターンを全通り列挙することだが、分解のパラメータは異なっていても最終的に同じ数列になるケースがあるので、それらを重複して数えないようにしないといけない。
そこで、分解の列挙順を定め、「より手前の分解順で得られる数列のパターンは除外してカウントする」というようにすると良い。

int N;
int P[33];
map<vector<int>, ll> memo;
const ll mo=1000000007;

vector<int> myslice(vector<int>& V,int a,int b) {
	return vector<int>(V.begin()+a,V.begin()+b);
}

ll hoge(vector<int> P) {
	if(P.size()<=1) return 1;
	//圧縮
	vector<int> Q=P;
	sort(ALL(Q));
	FORR(p,P) p=lower_bound(ALL(Q),p)-Q.begin();
	if(memo.count(P)) return memo[P];
	
	
	ll ret=0;
	int N=P.size();
	vector<ll> Xpat(N),Ypat(N);
	vector<int> R(N);
	int i,j,k;

	FOR(i,N) R[P[i]]=i;

	for(i=1;i<N;i++) {
		Xpat[i]=hoge(myslice(P,0,i));
		for(j=1;j<i;j++) Xpat[i]+=mo-Xpat[j]*hoge(myslice(P,j,i))%mo;
		int mi=*min_element(P.begin()+i,P.end());
		for(j=1;j<=mi;j++) {
			vector<int> W;
			FOR(k,i) if(P[k]>=j) W.push_back(P[k]);
			Xpat[i]+=mo-Ypat[j]*hoge(W);
		}
		Xpat[i]%=mo;
		ret+=Xpat[i]*hoge(myslice(P,i,N))%mo;
		
		Ypat[i]=hoge(myslice(R,0,i));
		
		for(j=1;j<i;j++) Ypat[i]+=mo-Ypat[j]*hoge(myslice(R,j,i))%mo;
		mi=*min_element(R.begin()+i,R.end());
		for(j=1;j<=mi;j++) {
			vector<int> W;
			FOR(k,i) if(R[k]>=j) W.push_back(R[k]);
			Ypat[i]+=mo-Xpat[j]*hoge(W);
		}
		Ypat[i]%=mo;
		ret+=Ypat[i]*hoge(myslice(R,i,N))%mo;
		
	}
	
	return memo[P]=ret%mo;
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N;
	vector<int> P(N);
	FOR(i,N) cin>>P[i];
	cout<<hoge(P)<<endl;
}

まとめ

こう書くとアプローチ自体はさほど珍しいものではないのだけど、分解の軸が2つあったりするのでややこしい。