AtCoder ABC #436 : G - Sum of Min

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

問題

N要素の整数列A、M要素の整数列B、正整数Kが与えられる。
を998244353で割った値を求めよ。

解法

A,BをあらかじめGCD(|A|,|B|)個に分割しておき、あとはGCD(|A|,|B|)=1の状態で考える。
A[i]が解に計上されるケースを考えると、A[i]がB[i],B[(i+H)%W],B[(i+2H)%W],B[(i+3H)%W]....より小さい場合である。
よって、A,Bの各要素を小さい順に走査し、B[i],B[(i+H)%W],B[(i+2H)%W],B[(i+3H)%W]....のうちすでにA[i]より小さいものの数をBITなどで数え上げて行けばよい。

int N,M;
ll A[202020],B[202020];
ll K;
const ll mo=998244353;

template<class V, int ME> class BIT {
public:
	V bit[1<<ME];
	V operator()(int e) {if(e<0) return 0;V s=0;e++;while(e) s+=bit[e-1],e-=e&-e; return s;}
	void add(int e,V v) { e++; while(e<=1<<ME) bit[e-1]+=v,e+=e&-e;}
};
BIT<int,20> X,Y;

ll hoge(vector<ll> A,vector<ll> B,ll K) {
	ll N=A.size(),M=B.size();
	ll loop=1LL*N*M;
	vector<pair<ll,ll>> V;
	int i;
	FOR(i,N) V.push_back({A[i],i});
	FOR(i,M) V.push_back({B[i],N+i});
	ll H=N,W=M;
	ll ret=0;
	sort(ALL(V));
	FORR2(a,b,V) {
		if(b<N) {
			(ret+=a*W)%=mo;
			H--;
		}
		else {
			(ret+=a*H)%=mo;
			W--;
		}
	}
	(ret*=K/loop%mo)%=mo;
	K%=loop;
	
	H=N,W=M;
	
	int RW,RH;
	for(i=1;i<W;i++) if(i*H%W==1) RW=i;
	for(i=1;i<H;i++) if(i*W%H==1) RH=i;
	FOR(i,2*W) X.add(i,1);
	FOR(i,2*H) Y.add(i,1);
	
	FORR2(a,b,V) {
		if(b<H) {
			ll num=K/H+(b<K%H);
			int start=b*RW%W;
			int len=X(start+num-1)-X(start-1);
			(ret+=1LL*a*len)%=mo;
			
			Y.add(b*RH%H,-1);
			Y.add(b*RH%H+H,-1);
		}
		else {
			b-=H;
			ll num=K/W+(b<K%W);
			int start=b*RH%H;
			int len=Y(start+num-1)-Y(start-1);
			(ret+=1LL*a*len)%=mo;

			X.add(b*RW%W,-1);
			X.add(b*RW%W+W,-1);
		}
	}
	
	
	return ret;
	
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>M>>K;
	FOR(i,N) cin>>A[i];
	FOR(i,M) cin>>B[i];
	int g=__gcd(N,M);
	ll ret=0;
	FOR(i,g) {
		vector<ll> As,Bs;
		FOR(j,N/g) As.push_back(A[j*g+i]);
		FOR(j,M/g) Bs.push_back(B[j*g+i]);
		ret+=hoge(As,Bs,K/g+(i<K%g));
	}
	cout<<ret%mo<<endl;
}

まとめ

なんとか思いつけて良かった。