~/leetcode$cat3524. Find X Value of Array I

Difficult ahh problem description to understand…

Walking through the example: [1, 2, 3, 4, 5] when k = 3 returns [9, 2, 4]. Let’s consider the 9 ways that are possible for x = 0, the example tells us its just all the possible ways to do the operations that do not remove nums[2] == 3. Therefore:

1. [1, 2, 3, 4, 5] - don’t remove
2. [2, 3, 4, 5] - remove 1
3. [1, 2, 3, 4] - remove 5
4. [3, 4, 5] - remove 1, 2
5. [2, 3, 4] - remove 1, 5
6. [1, 2, 3] - remove 4, 5
7. [3, 4] - remove 1, 2, 5
8. [2, 3] - remove 1, 4, 5
9. [3] - remove 1, 2, 4, 5

This seems awfully difficult for a medium problem… The first hint is lwk useless since this does seem like a dp problem instinctively. The second hint gives us an idea: Define dp[i][r] as the count of subarrays ending at index i whose product modulo k equals r. It does make sense that we have to count the number of subarrays ENDING at index i since we can only remove prefixes and suffixes (we can’t delete things from the middle of the array).

For the first example:

dp = [
    [0, 1, 0], (nums[i] = 1)
    [0, 0, 2], (nums[i] = 2)
    [3, 0, 0], (nums[i] = 3)
    [3, 1, 0], (nums[i] = 4)
    [3, 0, 2], (nums[i] = 5)
]

The update function isn’t too obvious on what it should be. Actually from what I see, there is a pattern: dp[i + 1][(nums[i + 1] * r) mod k] = dp[i][r] and then we have dp[i + 1][nums[i + 1] mod k] += 1. I think this sounds correct, and then at the end, we can just sum all the columns to get our final result array of length k.

Failed the second test case, and after printing dp, I realised the update has to be dp[i + 1][...] += dp[i][r] instead of just assigning.

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

        for i in range(n - 1):
            for r in range(k):
                dp[i + 1][(nums[i + 1] * r) % k] += dp[i][r]
            dp[i + 1][nums[i + 1] % k] += 1

        res = [0] * k
        for i in range(n):
            for r in range(k):
                res[r] += dp[i][r]

        return res

Time Complexity: O(n + k)

Space Complexity: O(n + k)

Time Taken: 34m 41s

Runtime: 583 ms

Memory: 48.1 MB

Tests Passed: 782/782