AtCoder ARC #149 : D - Simultaneous Sugoroku

どうにか解けて良かったね。
https://atcoder.jp/contests/arc149/tasks/arc149_d

問題

1次元の数直線上にN個の駒があり、それぞれの初期位置が与えられる。
それぞれ以下のクエリが計M個与えられる。

• 整数値Dが与えられる。駒を原点の向きにDだけ移動する。ただし駒がすでに原点にあるならそれ以上移動しない。

各駒が原点に停止するのは何回目の移動か。
もし原点に停止しない場合、最後の位置はどこか。

解法

登場し得る座標は-10^6～10^6である。
そこで、初期状態で[1,1000000]の区間を、平行移動しよう。
また、負の座標に入ってしまった区間は原点を中心に折り返そう。
その際、折り返して重なった部分はUnion-Findの要領でマージしていく。
そして初期状態の各位置が、いつ原点に行くのか、また最後どの位置に行くのかを計算していこう。

int N,M;
int X[303030],D[303030];

int ret[1010101];
int pos[1010101];

template<int um> class UF {
	public:
	vector<int> par,rank,cnt;
	UF() {par=rank=vector<int>(um,0); cnt=vector<int>(um,1); for(int i=0;i<um;i++) par[i]=i;}
	void reinit(int num=um) {int i; FOR(i,num) rank[i]=0,cnt[i]=1,par[i]=i;}
	int operator[](int x) {return (par[x]==x)?(x):(par[x] = operator[](par[x]));}
	int count(int x) { return cnt[operator[](x)];}
	int operator()(int x,int y) {
		if((x=operator[](x))==(y=operator[](y))) return x;
		cnt[y]=cnt[x]=cnt[x]+cnt[y];
		if(rank[x]>rank[y]) return par[x]=y;
		rank[x]+=rank[x]==rank[y]; return par[y]=x;
	}
};
UF<1<<20> uf;

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	
	cin>>N>>M;
	FOR(i,N) cin>>X[i];
	FOR(i,M) cin>>D[i];
	vector<pair<pair<int,int>,int>> S={{{1,1000000},1}};
	FOR(i,M) {
		vector<pair<pair<int,int>,int>> T;
		FORR2(p,f,S) {
			if(p.first>0) {
				p.first-=D[i];
				p.second-=D[i];
			}
			else {
				p.first+=D[i];
				p.second+=D[i];
			}
			if(p.first>0||p.second<0) {
				T.push_back({p,f});
			}
			else {
				if(p.first<0) T.push_back({{p.first,-1},f});
				if(p.second>0) T.push_back({{1,p.second},f+(1-p.first)});
				x=f-p.first;
				ret[uf[x]]=i+1;
			}
		}
		sort(ALL(T));
		S.clear();
		int pre=0;
		
		FORR(v,T) {
			int L=v.first.first,R=v.first.second,F=v.second;
			x=lower_bound(ALL(S),v)-S.begin();
			if(x&&S[x].first.first>L) x--;
			while(L<=R) {
				if(x==S.size()) break;
				if(L>S[x].first.second) {
					x++;
					continue;
				}
				uf(S[x].second+L-S[x].first.first,F);
				L++;
				F++;
			}
			if(L<=R) {
				S.push_back({{L,R},F});
			}
		}
	}
	FORR2(a,b,S) {
		for(i=a.first;i<=a.second;i++) pos[uf[b+i-a.first]]=i;
	}
	
	FOR(i,N) {
		x=uf[X[i]];
		if(ret[x]) {
			cout<<"Yes "<<ret[x]<<endl;
		}
		else {
			cout<<"No "<<pos[x]<<endl;
		}
			
	}
	
}

まとめ

今思えばもう少しすっきり書けるかもな。