Skip to content

3524. Find X Value of Array I

  • Time: $O(n^2)$
  • Space: $O(n)$
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
class Solution {
 public:
  vector<long long> resultArray(vector<int>& nums, int k) {
    vector<long long> ans(k);
    // dp[r] := the number of subarrays ending at current position with
    // product % k == r
    vector<long long> dp(k);

    for (const int num : nums) {
      vector<long long> newDp(k);
      const int numMod = num % k;
      // Start new subarray with only `num`.
      newDp[numMod] = 1;
      // Extend all previous subarrays.
      for (int i = 0; i < k; ++i) {
        const int newMod = (static_cast<long>(i) * numMod) % k;
        newDp[newMod] += dp[i];
      }
      // Accumulate counts into ans.
      for (int i = 0; i < k; ++i)
        ans[i] += newDp[i];
      dp = std::move(newDp);
    }

    return ans;
  }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
class Solution {
  public long[] resultArray(int[] nums, int k) {
    long[] ans = new long[k];
    // dp[r] := the number of subarrays ending at current position with
    // product % k == r
    long[] dp = new long[k];

    for (final int num : nums) {
      long[] newDp = new long[k];
      final int numMod = num % k;
      // Start new subarray with only `num`.
      newDp[numMod] = 1;
      // Extend all previous subarrays.
      for (int i = 0; i < k; ++i) {
        final int newMod = (int) (1L * i * numMod % k);
        newDp[newMod] += dp[i];
      }
      // Accumulate counts into ans.
      for (int i = 0; i < k; ++i)
        ans[i] += newDp[i];
      dp = newDp;
    }

    return ans;
  }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
class Solution:
  def resultArray(self, nums: list[int], k: int) -> list[int]:
    ans = [0] * k
    # dp[r] := the number of subarrays ending at current position with
    # product % k == r
    dp = [0] * k

    for num in nums:
      newDp = [0] * k
      numMod = num % k
      # Start new subarray with only `num`
      newDp[numMod] = 1
      # Extend all previous subarrays
      for i in range(k):
        newMod = (i * numMod) % k
        newDp[newMod] += dp[i]
      # Accumulate counts into ans
      for i in range(k):
        ans[i] += newDp[i]
      dp = newDp

    return ans