3525. Find X Value of Array II

Array
Math
Segment Tree

TIME

O((n + q) k log n)

SPACE

O(n * k)

Intuition

After dropping the prefix, the remaining choices are the non-empty prefixes of nums[start..]. So each query asks: among prefixes of a suffix of the array, how many have product ≡ x (mod k)? With point updates mixed in, that calls for a structure whose segments can be merged.

The information a segment needs is small because k <= 5: its total product mod k, and a histogram of the residues of its own prefix products. Two adjacent segments merge cleanly. The left segment's prefixes are unchanged. Each prefix of the right segment becomes a prefix of the combined segment after being multiplied by the left segment's full product, which just permutes (or collapses) its residue.

Approach

Segment tree over nums, each node holding prod and cnt[0..k-1].

Leaf for value v: prod = v % k, cnt[prod] = 1.
Merge (L, R): prod = L.prod * R.prod % k; cnt = copy of L.cnt; for each residue x, cnt[L.prod * x % k] += R.cnt[x].
Update index -> value: rewrite the leaf, recompute ancestors.
Query [start, n-1]: fold the covering nodes left to right with the same merge (order matters since multiplication by L.prod is applied to the right part only), then answer cnt[x].

The iterative bottom-up tree keeps a left accumulator and a right accumulator so the fold stays ordered without recursion.

Complexity

Time: O((n + q) * k * log n), build is O(n * k), each update and query touches O(log n) nodes with an O(k) merge
Space: O(n * k) for the tree

Edge Cases

k = 1 → every prefix has residue 0, answer is n - start when x = 0
start = n - 1 → only the single-element choice remains
Updates that do not change the value still run through the same path; no special case needed
Values up to 10^9 are reduced mod k at the leaf so products never grow
A residue of 0 anywhere in the retained prefix pins all longer prefixes to 0; the merge handles this automatically via L.prod = 0
Counts fit in int (at most n = 10^5 per query), unlike the sister problem which counts all subarrays

Code

from typing import List


class Solution:
    def resultArray(self, nums: List[int], k: int, queries: List[List[int]]) -> List[int]:
        """
        3525. Find X Value of Array II
        Time: O((n + q) * k * log n)
        Space: O(n * k)
        """
        n = len(nums)
        size = 1
        while size < n:
            size *= 2

        # node = (prod % k, cnt[x] = #non-empty prefixes with product % k == x)
        # identity: prod 1, empty histogram
        prod = [1] * (2 * size)
        cnt = [[0] * k for _ in range(2 * size)]

        def merge(pl, cl, pr, cr):
            out = cl[:]
            for x in range(k):
                if cr[x]:
                    out[pl * x % k] += cr[x]
            return pl * pr % k, out

        for i, v in enumerate(nums):
            m = v % k
            prod[size + i] = m
            cnt[size + i][m] = 1
        for node in range(size - 1, 0, -1):
            prod[node], cnt[node] = merge(
                prod[2 * node], cnt[2 * node], prod[2 * node + 1], cnt[2 * node + 1]
            )

        def update(pos, v):
            node = size + pos
            m = v % k
            prod[node] = m
            cnt[node] = [0] * k
            cnt[node][m] = 1
            node //= 2
            while node:
                prod[node], cnt[node] = merge(
                    prod[2 * node], cnt[2 * node], prod[2 * node + 1], cnt[2 * node + 1]
                )
                node //= 2

        def query(lo, hi):
            # inclusive [lo, hi]; fold left-to-right using two accumulators
            lp, lc = 1, [0] * k
            rp, rc = 1, [0] * k
            lo += size
            hi += size + 1
            while lo < hi:
                if lo & 1:
                    lp, lc = merge(lp, lc, prod[lo], cnt[lo])
                    lo += 1
                if hi & 1:
                    hi -= 1
                    rp, rc = merge(prod[hi], cnt[hi], rp, rc)
                lo //= 2
                hi //= 2
            return merge(lp, lc, rp, rc)[1]

        result = []
        for index, value, start, x in queries:
            update(index, value)
            result.append(query(start, n - 1)[x])
        return result

/**
 * 3525. Find X Value of Array II
 * Time: O((n + q) * k * log n)
 * Space: O(n * k)
 */
class Solution {
    private int k, size;
    private int[] prod;    // product of the node's segment mod k
    private int[][] cnt;   // cnt[node][x] = #non-empty prefixes with product % k == x

    public int[] resultArray(int[] nums, int k, int[][] queries) {
        this.k = k;
        int n = nums.length;
        size = 1;
        while (size < n) {
            size <<= 1;
        }
        prod = new int[2 * size];
        cnt = new int[2 * size][k];
        java.util.Arrays.fill(prod, 1);  // identity node: prod 1, empty histogram

        for (int i = 0; i < n; i++) {
            int m = nums[i] % k;
            prod[size + i] = m;
            cnt[size + i][m] = 1;
        }
        for (int node = size - 1; node >= 1; node--) {
            pull(node);
        }

        int[] result = new int[queries.length];
        for (int q = 0; q < queries.length; q++) {
            int index = queries[q][0], value = queries[q][1];
            int start = queries[q][2], x = queries[q][3];
            update(index, value);
            result[q] = query(start, n - 1)[x];
        }
        return result;
    }

    /** Recompute node from its two children (left then right). */
    private void pull(int node) {
        int l = 2 * node, r = 2 * node + 1;
        prod[node] = prod[l] * prod[r] % k;
        int[] out = cnt[node];
        System.arraycopy(cnt[l], 0, out, 0, k);
        for (int x = 0; x < k; x++) {
            out[prod[l] * x % k] += cnt[r][x];
        }
    }

    private void update(int pos, int value) {
        int node = size + pos;
        int m = value % k;
        prod[node] = m;
        java.util.Arrays.fill(cnt[node], 0);
        cnt[node][m] = 1;
        for (node >>= 1; node >= 1; node >>= 1) {
            pull(node);
        }
    }

    /** Histogram of prefix residues over nums[lo..hi], folded left to right. */
    private int[] query(int lo, int hi) {
        int lp = 1, rp = 1;
        int[] lc = new int[k], rc = new int[k];
        lo += size;
        hi += size + 1;
        while (lo < hi) {
            if ((lo & 1) == 1) {
                // acc = merge(acc, node)
                for (int x = 0; x < k; x++) {
                    lc[lp * x % k] += cnt[lo][x];
                }
                lp = lp * prod[lo] % k;
                lo++;
            }
            if ((hi & 1) == 1) {
                hi--;
                // acc = merge(node, acc)
                int[] merged = new int[k];
                System.arraycopy(cnt[hi], 0, merged, 0, k);
                for (int x = 0; x < k; x++) {
                    merged[prod[hi] * x % k] += rc[x];
                }
                rc = merged;
                rp = prod[hi] * rp % k;
            }
            lo >>= 1;
            hi >>= 1;
        }
        // final = merge(left acc, right acc)
        for (int x = 0; x < k; x++) {
            lc[lp * x % k] += rc[x];
        }
        return lc;
    }
}