LeetCode 3525 Elixir solution

LeetCode 3524 problem link: 3525. Find X Value of Array II

This is Hard version of the 3524. Find X Value of Array I problem.

This solution uses the Segment Tree, Prefix Product, and Modular Arithmetic patterns to handle permanent updates and range queries efficiently.
Each segment tree node stores the product of its entire segment modulo k and a tuple with the number of non-empty prefixes for each possible remainder. When two nodes are merged, prefixes from the left segment remain unchanged, while prefixes from the right segment are multiplied by the full product of the left segment. For every query, the value at nums[index] is permanently updated and only the nodes along the corresponding segment tree path are rebuilt.
After the update, the range [start, n - 1] is queried because removing the required prefix leaves exactly this part of the array. Each possible suffix removal then corresponds to a non-empty prefix of this range, so the count stored for remainder x gives the required x-value. Since k <= 5, each merge processes at most five remainders, making k effectively constant.



Time complexity: O(n + q * log n)
Space complexity: O(n)

where q is the number of queries.


Elixir

defmodule Solution do
  @spec result_array(nums :: [integer], k :: integer, queries :: [[integer]]) :: [integer]
  def result_array(nums, k, queries) do
    nums = List.to_tuple(nums)
    n = tuple_size(nums)

    # Build the segment tree in the process dictionary.
    build(nums, k, 1, 0, n - 1)

    # Apply each permanent update and query [start, n - 1].
    result = Enum.reduce(queries, [], fn [index, value, start, x], result ->
        update(k, 1, 0, n - 1, index, rem(value, k))
        {_prod, cnt} = query(k, 1, 0, n - 1, start, n - 1)
        [elem(cnt, x) | result]
      end)

    # Remove the stored segment tree nodes.
    cleanup(1, 0, n - 1)

    Enum.reverse(result)
  end

  # Build the segment tree.
  # A leaf has exactly one non-empty prefix: the element itself.
  defp build(nums, k, node, l, r) do
    if l == r do
      value = rem(elem(nums, l), k)
      Process.put(node, leaf(value))
    else
      mid = div(l + r, 2)

      build(nums, k, node * 2, l, mid)
      build(nums, k, node * 2 + 1, mid + 1, r)

      left = Process.get(node * 2)
      right = Process.get(node * 2 + 1)

      Process.put(node, merge_nodes(left, right, k))
    end
  end

  # Permanently update one element.
  # Only nodes on the path from the leaf to the root are rebuilt.
  defp update(k, node, l, r, index, value) do
    if l == r do
      Process.put(node, leaf(value))
    else
      mid = div(l + r, 2)

      if index <= mid do
        update(k, node * 2, l, mid, index, value)
      else
        update(k, node * 2 + 1, mid + 1, r, index, value)
      end

      left = Process.get(node * 2)
      right = Process.get(node * 2 + 1)

      Process.put(node, merge_nodes(left, right, k))
    end
  end

  # Return the combined information for range [ql, qr].
  defp query(k, node, l, r, ql, qr) do
    if ql <= l and r <= qr do
      Process.get(node)
    else
      mid = div(l + r, 2)

      cond do
        qr <= mid -> query(k, node * 2, l, mid, ql, qr)

        ql > mid -> query(k, node * 2 + 1, mid + 1, r, ql, qr)

        true -> 
            left = query(k, node * 2, l, mid, ql, qr)
            right = query(k, node * 2 + 1, mid + 1, r, ql, qr)

        merge_nodes(left, right, k)
      end
    end
  end

  # Merge two adjacent segments.
  # Left prefixes stay unchanged.
  # Right prefixes are multiplied by the full product of the left segment.
  defp merge_nodes({left_prod, left_cnt}, {right_prod, right_cnt}, k) do
    cnt = merge_counts(0, k, left_prod, left_cnt, right_cnt)
    {
      rem(left_prod * right_prod, k),
      cnt
    }
  end

  # Add all right prefix counts to their new remainders.
  defp merge_counts(r, k, _left_prod, cnt, _right_cnt) when r == k do
    cnt
  end

  defp merge_counts(r, k, left_prod, cnt, right_cnt) do
    nr = rem(left_prod * r, k)

    cnt = put_elem(cnt, nr, elem(cnt, nr) + elem(right_cnt, r))

    merge_counts(r + 1, k, left_prod, cnt, right_cnt)
  end

  # Create a leaf node with one non-empty prefix.
  defp leaf(value) do
    cnt = {0, 0, 0, 0, 0}

    {
      value,
      put_elem(cnt, value, 1)
    }
  end

  # Remove all tree nodes from the process dictionary.
  defp cleanup(node, l, r) do
    Process.delete(node)

    if l != r do
      mid = div(l + r, 2)

      cleanup(node * 2, l, mid)
      cleanup(node * 2 + 1, mid + 1, r)
    end
  end
end