[LeetCode] 3036. Number of Subarrays That Match a Pattern II

🔈 문제

You are given a 0-indexed integer array nums of size n, and a 0-indexed integer array pattern of size m consisting of integers -1, 0, and 1.

A subarray nums[i..j] of size m + 1 is said to match the pattern if the following conditions hold for each element pattern[k]:

• nums[i + k + 1] > nums[i + k] if pattern[k] == 1.
• nums[i + k + 1] == nums[i + k] if pattern[k] == 0.
• nums[i + k + 1] < nums[i + k] if pattern[k] == -1.

Return the count of subarrays in nums that match the pattern.

 

💫 예시

Input: nums = [1,2,3,4,5,6], pattern = [1,1]
Output: 4
Explanation: The pattern [1,1] indicates that we are looking for strictly increasing subarrays of size 3. In the array nums, the subarrays [1,2,3], [2,3,4], [3,4,5], and [4,5,6] match this pattern.
Hence, there are 4 subarrays in nums that match the pattern.

🎱 제한

• 2 <= n == nums.length <= 106
• 1 <= nums[i] <= 109
• 1 <= m == pattern.length < n
• -1 <= pattern[i] <= 1

 

📚 문제 풀이

 

문제를 정리해보자면 nums 에서 pattern 의 등장 횟수를 찾는 문제이다. 

nums 랑 pattern 이랑 바로 매칭이 되도록 nums[i] < nums[i+1] 이면 1, nums[i] == nums[i+1] 이면 0, nums[i] > nums[i+1] 이면 -1 인 새로운 배열을 만든다음 이 배열에서 pattern의 등장 횟수는 KMP 알고리즘을 사용하여 빠르게 해결할 수 있다.

 

💻 코드

class Solution:
    def countMatchingSubarrays(self, nums: List[int], pattern: List[int]) -> int:
        n = len(nums)
        pre = []
 
 
        for i in range(n-1):
            if nums[i] < nums[i+1]:
                pre.append(1)
            elif nums[i] > nums[i+1]:
                pre.append(-1)
            else:
                pre.append(0)
        
        # kmp
        m = len(pattern)
        fail = [0]*m
        i = 0
 
        for j in range(1, m):
            while i > 0 and pattern[i] != pattern[j]:
                i = fail[i-1]
            
            if pattern[i] == pattern[j]:
                i += 1
                fail[j] = i
        
        res = 0
        i = 0
        for j in range(n-1):
            while i > 0 and pre[j] != pattern[i]:
                i = fail[i-1]
            
            if pattern[i] == pre[j]:
                i += 1
                if i == len(pattern):
                    res += 1
                    i = fail[i-1]
 
        return res

 

🔗 문제링크 

LeetCode - The World's Leading Online Programming Learning Platform