(Python3) - LeetCode (Medium) : 3524. Find X Value of Array I

https://leetcode.com/problems/find-x-value-of-array-i

Find X Value of Array I - LeetCode

다소 adhoc 스러운 dp 였습니다.

📕 풀이방법

📔 입력 및 초기화

📑 정답 배열 ans, dict자료구조 cur을 선언 후 적절히 초기화해줍니다.

📔 풀이과정

문제 요구사항을 분해해 재조립하면 답이 보이는 문제입니다. 문제의 연산은 prefix와 suffix를 각각 제거하고, 배열을 비어 있지 않게 남기는 것입니다.

 

예를 들어

nums = [1, 2, 3, 4]

에서 prefix [1], suffix [4]를 제거하면 [2, 3]이 남습니다.

이를 일반화한다면 어떤 방식으로 prefix와 suffix를 제거하더라도 최종적으로 남는 배열은 항상

nums[left ... right]

형태의 연속된 부분 배열입니다.

 

즉 문제를 다시 표현하면,

nums의 모든 비어 있지 않은 연속 부분 배열에 대해
부분 배열 원소의 곱 % k를 구하고, 각 나머지가 몇 번 나오는지 세는 문제입니다.

가 됩니다.

 

가장 단순하게는 모든 (left, right)를 탐색할 수 있습니다. 하지만 가능한 (left, right)의 수 자체가 O(N²)개이므로, memoization을 사용하더라도 각 구간을 상태로 저장한다면 O(N²)개의 상태를 줄일 수 없습니다.

따라서 (left, right) 자체를 저장하지 않고 더 작은 상태로 줄일 방법을 생각해야 합니다.

 

이때 모든 연속 부분 배열을 오른쪽 끝점 right가 같은 것끼리 묶어서 생각할 수 있습니다. right를 하나 고정하면 가능한 구간은 left만 다른 형태가 되고, right를 한 칸 증가시킬 때는 기존 구간의 오른쪽에 새 원소 하나를 붙이는 형태가 됩니다.

 

따라서 right를 왼쪽에서 오른쪽으로 하나씩 이동시키면서, 이전 right에서 계산한 정보를 재사용할 수 있는지 살펴봅니다.

 

📑 왜 right를 하나씩 이동시키는가
모든 부분 배열에는 정확히 하나의 오른쪽 끝점이 있습니다.

 

예를 들어

nums = [1, 2, 3, 4]

라면 오른쪽 끝점별로 가능한 배열은 다음과 같습니다.

right = 0
[1]

right = 1
[1, 2]
[2]

right = 2
[1, 2, 3]
[2, 3]
[3]

right = 3
[1, 2, 3, 4]
[2, 3, 4]
[3, 4]
[4]

이렇게 보면 모든 연속 부분 배열이 정확히 한 번씩 등장합니다.

 

원래 문제의 prefix/suffix 관점으로 보면,

right = 2

에서 세는 배열들은 모두 [4]를 suffix로 제거한 경우이고,

[1,2,3]   → prefix 제거 없음
[2,3]     → prefix [1] 제거
[3]       → prefix [1,2] 제거

에 해당합니다.

 

따라서 right를 0부터 nums길이-1까지 이동하면 suffix를 어디까지 제거할지가 자연스럽게 모두 처리되고, 각 right에서 여러 prefix제거 경우를 처리하면 됩니다.

 

📑 left를 모두 순회하지 않는 법

여기까지 그대로 구현하면 각 right마다 모든 left를 다시 확인해야 하므로 여전히 O(N²)입니다.

하지만 우리는 각 부분 배열의 실제 곱 전체가 필요한 것이 아니라 곱 % k만 필요합니다.

현재 위치에서 끝나는 두 부분 배열의 곱을 각각 A,B라고 할때

A % k == B % k

라면 다음 숫자 num을 두 배열에 추가했을 때도

(A * num) % k == (B * num) % k

입니다.

 

따라서 곱의 나머지가 같은 부분 배열들은 앞으로도 똑같이 변화하므로 서로 구분할 필요가 없습니다.

그래서 다음과 같이 줄일 수 있습니다.

cur[remainder] =
    현재 right에서 끝나는 부분 배열 중
    product % k == remainder 인 배열의 개수

 

예를 들어

cur = {2: 5}

라면 실제로 어떤 left에서 시작했는지는 저장하지 않고,

현재 위치에서 끝나며 곱의 나머지가 2인 부분 배열이 5개 있다.

는 정보만 저장합니다.

다음 숫자가 num이라면 이 5개는 각각 계산할 필요 없이 한꺼번에

new_remainder = (2*num)%k

의 위치에 해당하는 곳으로 이동시킬 수 있고, 이 위치에 5를 더하면 됩니다.

 

단, 기존 부분 배열을 연장하는 경우 외에도 현재 num하나만 남기는 경우가 새로 하나 생깁니다. 이는 현재 위치 이전의 모든 원소를 prefix로 제거한 경우이므로 num % k 위치에도 하나 추가합니다.

new_cur[num % k] += 1

 

 

이렇게 만들어진 next_cur은 현재 right에서 끝나는 모든 부분 배열을 나타냅니다. 따라서 각 (remainder, count)에 대해

ans[remainder] += count

를 수행합니다.

 

이후

cur = next_cur

로 갱신해 다음 right를 처리합니다.

 

📑 시간 복잡도

O(N * K): 각 nums에 대해 현재 존재하는 나머지 상태를 순회하며, 가능한 나머지는 0 ~ k-1이므로 최대 k개의 상태만 존재합니다. 문제에서 k <= 5이므로 k를 상수로 보면 최종적으로 O(N)입니다.

 

📑 공간 복잡도

O(k): cur, next_cur, ans가 각각 최대 k개의 상태만 저장합니다. 문제에서 k <= 5이므로 O(1)로 볼 수 있습니다.

📔 정답 출력 | 반환

ans를 반환합니다.

---

📕 Code

📔 Python3

class Solution:
    def resultArray(self, nums: List[int], k: int) -> List[int]:
        ans = [0] * k
        cur = {}
        for num in nums:
            next_cur = {}
            
            for item, count in cur.items():
                next_rem = item*num % k
                next_cur[next_rem] = next_cur.get(next_rem,0) + count
            
            next_rem = num % k
            next_cur[next_rem] = next_cur.get(next_rem,0) + 1

            for rem, cnt in next_cur.items():
                ans[rem] += cnt
            cur = next_cur
        return ans

---

*더 나은 내용을 위한 지적, 조언은 언제나 환영합니다.