kmjp's blog

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

LeetCode Weekly Contest 123 : 992. Subarrays with K Different Integers

3問目の方が手こずった。
https://leetcode.com/contest/weekly-contest-123/problems/subarrays-with-k-different-integers/

問題

整数列Aが与えられる。
この部分列のうち、数列内で登場する値がちょうどK種なのは何通りか。

解法

尺取り法の要領でK種以下となる部分列は容易に数えられる。
そこで、K種以下の部分列の数から、(K-1)種以下の部分列の数を引こう。

class Solution {
public:
    int subarraysWithKDistinct(vector<int>& A, int K) {
        int L1,L2,R,ret=0;
        
        L1=L2=0;
        map<int,int> M1,M2;
        FOR(R,A.size()) {
			M1[A[R]]++;
			M2[A[R]]++;
			while(M1.size()>K) {
				M1[A[L1]]--;
				if(M1[A[L1]]==0) M1.erase(A[L1]);
				L1++;
			}
			while(M2.size()>K-1) {
				M2[A[L2]]--;
				if(M2[A[L2]]==0) M2.erase(A[L2]);
				L2++;
			}
			ret+=L2-L1;
		}
		
        return ret;
    }
};

まとめ

なんか今回WA出しまくった。