3524. Find X Value of Array I

[zeikar]

Problem Link

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

Problem Summary

배열 앞뒤를 잘라내고 남은 부분의 곱을 k로 나눈 나머지별로 경우의 수를 세는 문제.

Solution

문제가 prefix, suffix를 떼는 식으로 되어 있어서 처음엔 잘 안 보였는데, 결국 남는 건 비어있지 않은 부분 배열 하나다. 즉 모든 부분 배열의 곱 % k를 세는 문제.

처음엔 prefix sum처럼 prefix product로 구간 곱을 꺼내려고 했는데 안 된다. 구간 곱을 꺼내려면 나눗셈이 필요한데 mod k에서는 역원이 없을 수 있다. (k = 4에서 2처럼)

여기서 k <= 5가 유독 작다. 곱 % k는 많아야 5종류니까, 끝나는 인덱스별로 부분 배열들을 나머지별 개수로 요약해서 들고 다니면 된다. cnt[i][r]을 i에서 끝나는 부분 배열 중 곱 % k == r인 개수라고 하면, 다음 원소 a가 들어올 때 나머지 r인 칸의 개수는 (r * a) % k 칸으로 옮겨가고 a 하나짜리 부분 배열이 a 칸에 하나 더해진다. 모든 부분 배열은 끝나는 인덱스가 딱 하나니까 끝점별 cnt를 전부 더해주면 답이다.

즉 DP인데, 부분 배열 세기로 포장돼 있어서 생각해내기가 쉽지 않았다... 결국 claude의 도움을 받아서 풀었다.

시간복잡도는 O(nk).

Source Code

class Solution:
    def resultArray(self, nums: List[int], k: int) -> List[int]:
        n = len(nums)
        for i in range(n):
            nums[i] %= k

        ans = [0 for _ in range(k)]
        cnt = [[0 for _ in range(k)] for _ in range(n+1)]
        
        for i in range(n):
            for j in range(k):
                cnt[i+1][j*nums[i]%k] += cnt[i][j]
            cnt[i+1][nums[i]] += 1

            for j in range(k):
                ans[j] += cnt[i+1][j]

        return ans