Computer Science/Algorithm

[Python 3] LEETCODE 2962. Count Subarrays Where Max Element Appears at Least K Times

무니화니 2024. 3. 29. 16:31

 

한국말로 번역하자면,

"당신은 정수로 이루어진 배열인 nums와 양의 정수 k 가 주어진다.

부분 수열 중 nums의 제일 큰 값이 최소 k번 이상 위치한 부분 수열의 개수를 return하라

여기서 부분 수열이란 배열 내 연속적인 원소들의 나열이다."

 

여기서 슬라이딩 윈도우 알고리즘을 사용하여 문제를 해결하였다.

슬라이딩 윈도우 알고리즘이란, 투 포인터 알고리즘의 일부이다.

슬라이딩 원도우는 연속되어있는 두개의 포인터가 사용되어 특정 조건을 해결하는 알고리즘이다. 

여기서 정확히 말하면 슬라이딩 윈도우의 약간의 변형이다. 만약 해당 배열에서 제일 큰 원소의 값을 M이라고 하였을 때, 오른쪽 포인터가 M인 경우에는 슬라이딩 윈도우가 적용되어 왼쪽에 있는 포인터가 오른쪽으로 이동하지만, 만약 M이 아닌 경우에는 왼쪽 포인터가 이동하지 않는다.

 

class Solution:
    def countSubarrays(self, nums: List[int], k: int) -> int:
        M=max(nums)
        n=len(nums)
        cnt=0
        ans=0
        l=0
        for r in range(n):
            if nums[r]==M: 
                cnt+=1
            while cnt>=k:
                if nums[l]==M: 
                    cnt-=1
                l+=1
            ans+=l
        return ans

 

작성한 정답 코드이다. 되게 짧고 간단해보이지만, 직관적이지 않을 수 있다. Example 1을 그림을 통해서 설명하겠다.

 

 

 

각 인덱스를 0,1,2,3,4 라고 해보자.

포인터 Left와 Right이 있다. 현재 k의 값은 2이므로, 해당 list의 최댓값이 3이기에 두 개 이상의 3이 필요하다.

Right의 인덱스가 3일때, 드디어 2개의 3이 있다.

즉, 이런 부분집합이 정답에 포함되어야 할 것이다.

 

 

 

Left의 인덱스를 오른쪽으로 한 칸 옮겨보자. [3,2,3], 아직까지 충분히 가능하다. 이것도 정답에 포함되어야한다.

 

하지만 여기서 키 포인트는, 다음 Right을 인덱스 4로 옮겼을 때, Left를 인덱스 0으로 리셋해줄 필요가 없다는 것이다. 왜냐하면, 이미 Right의 인덱스가 3일때부터도 Left의 인덱스가 0일때, 1일때 모두 가능한 부분집합이였기 때문이다.

즉, Right의 인덱스가 4일때, Left의 인덱스가 0,1 일때, 즉 [1,3,2,3,3]과 [3,2,3,3]은 모두 2개 이상의 3이 존재한다.

여기서 사용되는 아이디어가 슬라이딩 윈도우의 변형이라고 할 수 있겠다.

 

L을 다시 0으로 보내서 확인할 필요가 없다. R을 4로 옮겨도, 1부터 확인하자.

 

이런 식으로, L은 3까지 전진해도 [3,3]이라는 부분집합은 2개의 3을 가지며 조건을 만족한다.

즉, R이 3일때 L이 0,1, R이 4일때 L이 0,1,2,3인 경우, 모두 6가지가 가능한 부분집합의 케이스이다.