TopCoder SRM 818 : Div1 Easy Div2 Hard MedianSegments

せっかく本番余裕をもって全完できたのに、unratedだった…。
https://community.topcoder.com/stat?c=problem_statement&pm=17274&rd=18920

問題

整数列Aと整数Kが与えられる。
A[K]はAの中で1回しか現れない。
K番目の要素を含むAの部分列A[L...R]を取ったとき、奇数長かつmedianがA[K]となるのは何通りか。

解法

Aの各要素を

• A[K]未満のものを-1
• A[K]を0
• A[K]より大きいものを1

と置き換えたとき、部分列の総和が0になれば条件を満たす。

そこで、左端Lを走査したときのA[L...(K-1)]の総和と右端Rを走査したときのA[(K+1)...R]の総和で、値が正負反転したものになる組み合わせを数え上げよう。

class MedianSegments {
	public:
	long long count(int N, int K, vector <int> Aprefix, int M) {
		vector<ll> A;
		FORR(a,Aprefix) A.push_back(a);
		ll state=A.back();
		int i;
		for(i=Aprefix.size();i<N;i++) {
			state = (state * 1103515245 + 12345)%(1LL<<31);
			A.push_back(state%M);
		}
		vector<ll> V;
		
		ll ret=0;
		map<int,int> T;
		int cur=0;
		T[0]=1;
		for(i=K+1;i<N;i++) {
			if(A[i]>A[K]) cur--;
			else if(A[i]<A[K]) cur++;
			T[cur]++;
		}
		cur=0;
		ret+=T[0];
		for(i=K-1;i>=0;i--) {
			if(A[i]>A[K]) cur++;
			if(A[i]<A[K]) cur--;
			ret+=T[cur];
		}
		return ret;
	}
}

まとめ

解法よりも入力値の生成に手間取るので、この入力形式はどうにかしてくれないかな。