kmjp's blog

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

LeetCode Weekly Contest 107 : 927. Three Equal Parts

本番参加は検討中。ひとまず難易度8以上の問題を様子見ていきます。
https://leetcode.com/contest/weekly-contest-107/problems/three-equal-parts/

問題

01で構成される数列がある。
これを3つの連続部分列に分割し、それぞれを2進数表記とみなしたとき同じ値となるようにできるか。

解法

まず1の数を数えると、各部分列に含むべき1の数が確定する。
次に、末尾の0の数を数えると、各部分列で最後の1の後につく0の数が確定する。

これで分割位置が確定するので、leading zeroを除いて一致するか判定すればよい。

class Solution {
public:
	vector<int> threeEqualParts(vector<int>& A) {
		int i;
		int N=A.size();
		int N1=count(ALL(A),1);
		if(N1%3) return {-1,-1};
		if(N1==0) return {0,2};
		int L0=0;
		FOR(i,N) {
			if(A[N-1-i]) break;
			L0++;
		}
		int X,Y;
		int L1=0;
		FOR(i,N) if(A[i]) {
			L1++;
			if(L1==N1/3) X=i+L0;
			if(L1==N1*2/3) Y=i+L0;
		}
		
		vector<int> V[3];
		FOR(i,N) {
			int id=(i<=X)?0:((i<=Y)?1:2);
			if(A[i] || V[id].size()) V[id].push_back(A[i]);
		}
		
		if(V[0]!=V[1]) return {-1,-1};
		if(V[0]!=V[2]) return {-1,-1};
		if(V[1]!=V[2]) return {-1,-1};
		return {X,Y+1};
		
		
	}
};

まとめ

HackerRankの80ptとは全然違う難易度だ。