kmjp's blog

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

AtCoder ARC #054 : D - バブルソート

ARC

これはライブラリも持ってなかったしなぁ…。
http://arc054.contest.atcoder.jp/tasks/arc054_d

問題

以下のステップで作られる巨大な数列を考える。
入力を処理すると、スタック内には単一の数列が残る。
この数列に対しバブルソートを実行した際の要素のswap回数を求めよ。

一時領域として数列のスタックがあり、またN個の整数A[i]を順に処理することで数列を生成する。

  • A[i]が正なら、A[i]1要素だけの数列をスタックにpushする。
  • A[i]が0なら、スタックの上２つの要素をpopし、上から2つめの数列に1つ目の数列を連結したものを再度スタックにpushする。
  • A[i]が負なら、スタックの最上位の数列を、A[i]回繰り返したものに置き換える。

解法

Editorialを見て回答。

まず部分点を考える。
スタック内の各数列Sについて、律儀に数列本体を持つとMLEするので、以下の３値だけを持つことにする。

  • Sx := この数列をバブルソートした際のswap回数
  • Sy := この数列を２回繰り返したとき、バブルソートした際のswap回数が2Xより何回多いか
  • Sd := 数列内の各要素の登場回数

上記３通りの数列生成手順に対し、Sx,Sy,Sdを適切に変化させよう。
最終的に数列が最後1個残った場合、Sxが解。

  • A[i]が正の場合
    • Sx = Sy = 0, SdにはA[i]を1個いれた数列Sを生成してpush
  • A[i]が負の場合
    • k=-A[i]とする。スタックの最上位の数列Sを、Tx = k*Sx + k*(k-1)*Sy, Ty = k^2*Sy, Td = Sdの各要素をk倍したもの、である数列Tに置き換える。
  • A[i]が0の場合
    • スタックの最上位の数列T,２番目Sに対し、Ux = Sx + Tx, Uy = Sy + Ty + a , Ud = Sd + Tdで生成される数列Uに置き換える。
    • aはUのバブルソート時にS内の要素とT内の要素をswapする回数で、SdとTdから求められる。

満点解法では、最後のUyやUdを求めるを愚直にやるとTLEする。
Udはマージテクによりデータのコピー回数を削減し、SdとTdからaを求める手順はSegTreeを使って計算量を減らそう。
なお、SegTreeを使う場合は有効な要素が入っている箇所だけメモリを取るようにしないとMLEする。

ll mo=1000000007;

template<class V,int MV> struct Dyn_SegTree {
	struct IntNode {
		IntNode(){ val=0, left=right=NULL;}
		V val;
		IntNode *left,*right;
	};
	vector<IntNode*> pool;
	IntNode root;
	IntNode* getent() {
		pool.push_back(new IntNode);
		return pool.back();
	};
	~Dyn_SegTree(){
		FORR(r,pool) delete r;
	}
	
	void add(int k,ll v,int l,int r,IntNode* node) { //val[k]+=v
		if(k<l || k>=r) return;
		node->val += v;
		if(node->val>=mo) node->val -= mo;
		if(l+1==r) return;
		int m=(l+r)>>1;
		if(k>=l && k<m) {
			if(node->left==NULL) node->left=getent();
			add(k,v,l,m,node->left);
		}
		if(k>=m && k<r) {
			if(node->right==NULL) node->right=getent();
			add(k,v,m,r,node->right);
		}
	}
	void add(int k,V v){ add(k,v,0,MV,&root);}
	
	V sum(int k,int l,int r,IntNode* node) { //sum(val[0..(k-1)])
		if(l==r || k<l) return 0;
		if(k>=r) return node->val;
		V v=(node->left?sum(k,l,(l+r)>>1,node->left):0)+(node->right?sum(k,(l+r)>>1,r,node->right):0);
		if(v>=mo) v-=mo;
		return v;
	}
	//sum(val[0..k])
	V sum(int k){ return sum(k+1,0,MV,&root);}
};

struct Node {
	ll X,Y;
	ll mul;
	ll sum;
	unordered_map<int,ll> D;
	Dyn_SegTree<int,1<<17> ds;
};

list<Node> pool;
vector<Node*> S;
int N;

ll modpow(ll a, ll n = mo-2) {
	ll r=1;
	while(n) r=r*((n%2)?a:1)%mo,a=a*a%mo,n>>=1;
	return r;
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	
	cin>>N;
	FOR(i,N) {
		cin>>x;
		
		if(x>0) {
			pool.insert(pool.end(),Node());
			Node* n=&pool.back();
			n->X = n->Y = 0;
			n->mul=n->sum=n->D[x]=1;
			n->ds.add(x,1);
			S.push_back(n);
		}
		if(x==0) {
			pool.insert(pool.end(),Node());
			Node* n=&pool.back();
			Node* t=S.back();
			S.pop_back();
			Node* s=S.back();
			S.pop_back();
			
			n->X = (s->X + t->X)%mo;
			n->Y = (s->Y + t->Y)%mo;
			
			if(s->D.size()<t->D.size()) {
				FORR(sr,s->D) {
					ll small=t->ds.sum(sr.first-1)*t->mul%mo;
					ll same=t->D[sr.first]*t->mul%mo;
					
					(n->X += sr.second*s->mul%mo*(small%mo))%=mo;
					(n->Y += sr.second*s->mul%mo*((t->sum-same+mo)%mo))%=mo;
				}
			}
			else {
				FORR(sr,t->D) {
					ll small=s->ds.sum(sr.first-1)*s->mul%mo;
					ll same=s->D[sr.first]*s->mul%mo;
					ll large=((s->sum-small-same)%mo+mo)%mo;
					
					(n->X += sr.second*t->mul%mo*(large%mo))%=mo;
					(n->Y += sr.second*t->mul%mo*((s->sum-same+mo)%mo))%=mo;
				}
				swap(s,t);
			}
			
			n->mul=t->mul;
			n->sum=s->sum+t->sum;
			swap(n->D,t->D);
			swap(n->ds.root,t->ds.root);
			swap(n->ds.pool,t->ds.pool);
			
			ll rev=modpow(n->mul);
			FORR(sr,s->D) {
				ll v=sr.second*s->mul%mo;
				(n->D[sr.first] += v*rev)%=mo;
				n->ds.add(sr.first,v*rev%mo);
			}
			
			S.push_back(n);
		}
		if(x<0) {
			Node* n=S.back();
			ll a=-x;
			n->X = (n->X*a+(a*(a-1)/2)%mo*n->Y)%mo;
			n->Y = a*a%mo*n->Y%mo;
			(n->mul *= a)%=mo;
			(n->sum *= a)%=mo;
		}
	}
	
	cout<<S[0]->X<<endl;
	
	
}

まとめ

動的に領域を確保するSegTree初めて書いた。