LeetCode Weekly Contest 409 : 3245. Alternating Groups III

LeetCodeの問題の割にコードが長め。
https://leetcode.com/problems/alternating-groups-iii/description/

問題

円環状にN個のタイルが並んでいる。
各タイルは赤か青で塗られている。

以下のクエリに順次答えよ。

• 1つのタイルの色を変える
• 整数Kが与えられる。長さKの連続するタイルで、赤青交互に並ぶようなものは何通りあるか。

解法

奇数番目のタイルの赤青を反転させておくと、結局同じ色がK個以上連続するタイルを数える問題となる。
以下はEditorialとは違う解法。

色が入れ替わる位置をsetで持ち、同色が並ぶ最長のタイル数と、その頻度をmapで持とう。
後者は高々O(√N)通りしかない。
なので、2つ目にクエリに対してはO(√N)通りに対し、K個の連続タイトルを取れる数を持てばよい。

円環状にタイルが並んでいる点だけ注意。

set<int> S;
map<int,int> M;
int N;

class Solution {
public:
	void update(int x) {
		if(x==0||x==2*N) return;
		if(S.count(x)) {
			auto it=S.lower_bound(x);
			int a=*(prev(it));
			int b=*(next(it));
			S.erase(x);
			if(x<N) {
				if(--M[x-a]==0) M.erase(x-a);
			}
			if(b<N) {
				if(--M[b-x]==0) M.erase(b-x);
			}
			if(b<N) M[b-a]++;
			
		}
		else {
			auto it=S.lower_bound(x);
			auto pre=prev(it);
			int a=*pre;
			int b=*it;
			if(b<N) {
				if(--M[b-a]==0) M.erase(b-a);
			}
			S.insert(x);
			if(x<N) M[x-a]++;
			if(b<N) M[b-x]++;
		}
		
	}
    vector<int> numberOfAlternatingGroups(vector<int>& colors, vector<vector<int>>& queries) {
        N=colors.size();
        int i;
        FOR(i,N) colors.push_back(colors[i]);
        S.clear();
        S.insert(0);
        S.insert(2*N);
        M.clear();
        FOR(i,2*N) {
			if(i%2) colors[i]^=1;
			if(i&&colors[i]!=colors[i-1]) update(i);
		}
		
		vector<int> ret;
		FORR(q,queries) {
			if(q[0]==1) {
				int sum=0;
				if(*next(S.begin())>N) {
					sum=N;
				}
				else {
					FORR(a,M) if(a.first>=q[1]) sum+=1LL*(a.first-q[1]+1)*a.second;
					auto it=S.lower_bound(N);
					int b=*it;
					int a=*prev(it);
					sum+=max(0,min(b-a-q[1]+1,N-a));
				}
				ret.push_back(sum);
				
			}
			else {
				if(colors[q[1]]!=(q[2]^(q[1]%2))) {
					colors[q[1]]^=1;
					update(q[1]);
					update(q[1]+1);
					update(q[1]+N);
					update(q[1]+1+N);
					
				}
			}
		}
		return ret;
        
    }
};

まとめ

LeetCodeで1クエリO(√N)掛けるの心配だったけど、どうにか間に合った。