[leetcode] Count Subarrays Where Max Element Appears at Least K Times

문제

• 링크
• 배열 nums와 자연수 k가 주어진다. nums에 있는 최대값을 M이라고 할 때, nums의 subarray 중에서 M이 k개 이상 포함되는 subarray 개수를 구해야 한다.
• 제약 조건
  • 배열 길이: 1 <= nums.length <= 10^5
  • 배열 값 범위: 1 <= nums[i] <= 10^6
  • k 범위: 1 <= k <= 10^5

풀이

풀기 전

• 함수의 반환 타입부터가 long이다. 모든 경우의 수를 파악해서는 안 될 거 같은 느낌이다. 실제로 배열 길이를 n이라고 하면 모든 subarray 개수는 n(n+1)/2이다. 제약 조건을 보면 최대 약 50억이다. 너무 크다.
• 시간을 줄이기 위해서 생각해보면, 조건을 만족하는 subarray가 있을 때 해당 subarray를 포함하는 subarray도 조건을 만족한다는 걸 알 수 있다. 그러면 일단 조건을 만족하는 최단 길이의 subarray들을 구해볼 수 있다.
  • 그럼 투 포인터 방식으로 접근할 수 있다. left는 조건을 만족하는 subarray의 가장 왼쪽 인덱스를 가리키고, right는 가장 오른쪽 인데스를 가리키게 하면 된다. 이때 left와 right 위치에 있는 값은 모두 최대값 M일 것이다.
  • 중복으로 집계하면 머리가 아프므로 right로 순회 중인 인덱스까지만 고려한다. 그 뒤에 순회하는 값들은 순회하는 시점에 고려한다. 예를 들어, nums = [1, 3, 3, 1], k = 2가 있을 때, right=2이면 [1, 3, 3]까지만 생각한다. 마지막 1은 right가 4일 때 집계한다.
  • left와 right 각각 한 번씩만 nums를 순회하므로 시간 복잡도는 O(nums.length)이다. 추가로 사용하는 배열이나 맵 등은 없으므로 공간 복잡도는 O(1)이다.

코드

class Solution {
    public long countSubarrays(int[] nums, int k) {
        long ans = 0;

		// 최대값을 찾는다.
        int maxVal = 0;
        for (Integer num : nums)
            maxVal = Math.max(maxVal, num);

        int num;
        int left = 0, right = 0;
        int maxFreq = 0;  // 최대값의 빈도수를 저장한다.
        for (; right<nums.length; right++) {
            num = nums[right];

			// 현재 값이 최대값이 아닌 경우,
            // 1) 빈도수가 k보다 작으면 그냥 넘어간다.
            // 2) 빈도수가 k보다 크거나 같으면 left+1을 해준다. [0, left-1] 범위에서 [0, left] 범위로 확장했을 때 추가되는 subarray 개수를 더해준다고 보면 된다.
            if (num != maxVal) {
                if (maxFreq < k)
                    continue;

                ans += left + 1;
            } else {  // 현재 값이 최대값인 경우이다.
                maxFreq++;  // 빈도수 1 올려준다.
                
                // 빈도수가 k보다 작으면 건너뛴다.
                if (maxFreq < k)
                    continue;
                
                // 빈도수가 k보다 크면 left를 하나 올리고 빈도수를 1 줄여준다.
                // 왜냐하면 조건을 만족하는 가장 작은 subarray을 찾기 위해서이다.
                // 정확히 하려면 nums[left]==maxVal 조건을 추가해야 하지만, 해당 코드를 진행하면 조건이 이미 전제되어 있어서 뺐다.
                if (maxFreq > k) {
                    left++;
                    maxFreq--;
                }

				// 조건을 만족하는 가장 작은 subarray를 찾기 위해 left를 최대값을 찾을 때까지 올린다.
                while (nums[left] != maxVal)
                    left++;

				// [0, left-1] 범위에서 [0, left] 범위로 확장했을 때 추가되는 subarray 개수를 더해준다고 보면 된다.
                ans += left + 1;
            }
        }
        return ans;
    }
}

푼 후

• 처음엔 복잡한가 싶었는데, 투 포인터 떠올리니 무난하게 풀었다.