LeetCode 3524 Elixir solution

LeetCode 3524 problem link: 3524. Find X Value of Array I

This solution uses Dynamic Programming and Modular Arithmetic technique to count all possible remaining subarrays. The dp[r] state stores the number of subarrays ending at the previous position whose product has remainder r modulo k.

For each number, we start a new subarray with that element and extend all previous subarrays by updating their remainders with (r * mod) % k. The counts in new_dp are then added to result, which stores the total number of subarrays for each possible remainder. Since k <= 5, we only maintain a small number of remainder states instead of calculating potentially huge products.

Time complexity: O(n * k)
Space complexity: O(k)

Elixir

defmodule Solution do
  @spec result_array(nums :: [integer], k :: integer) :: [integer]
  def result_array(nums, k) do
    dp = List.duplicate(0, k)
    result = List.duplicate(0, k)

    {_, result} =
      Enum.reduce(nums, {dp, result}, fn num, {dp, result} -> mod = rem(num, k)

        new_dp =
          dp
          |> Enum.with_index()
          |> Enum.reduce(List.duplicate(0, k), fn {count, r}, acc ->
            new_r = rem(r * mod, k)
            List.update_at(acc, new_r, &(&1 + count))
          end)
          |> List.update_at(mod, &(&1 + 1))

        result = Enum.zip_with(result, new_dp, &(&1 + &2))

        {new_dp, result}
      end)

    result
  end
end