403. Find X Value of Array II

Because k <= 5, each segment-tree node can store, for every possible incoming remainder, how many prefixes inside its range land on each remainder mod k. Merging two nodes just feeds the left node's exit remainder into the right node.

Reframing the query

Strip away the story. After the persistent update nums[index] = value, a single query with parameters (start, x) asks:

Among the elements nums[start], nums[start+1], ..., nums[n-1], how many non-empty prefixes have a product that is congruent to x modulo k?

That is because "remove a suffix but keep the array non-empty" over the array nums[start..n-1] is exactly "keep a non-empty prefix of nums[start..n-1]". So we count indices j in [start, n-1] with (nums[start] * ... * nums[j]) mod k == x.

Why brute force is too slow

Each query is answered in O(n) by walking the running product from start. With n up to 10^5 and up to 2 * 10^4 queries that is 2 * 10^9 operations — too slow. We also have point updates between queries, so we need a structure that supports both updates and range aggregation.

The key observation: k is tiny (k <= 5)

Work with every value mod k. A running product only ever has one of k possible remainders (0, 1, ..., k-1). For a contiguous block of the array define, for each incoming remainder a:

• cnt[a][b] = the number of prefixes ending inside this block whose product (started from remainder a) equals remainder b.
• g[a] = the remainder you exit the whole block with, i.e. a * (product of the block) mod k.

A single element m = nums[i] mod k is a leaf: for every a, its one prefix lands on b = a*m mod k, so cnt[a][a*m mod k] = 1 and g[a] = a*m mod k.

Merging two blocks (L then R)

Enter L with remainder a. Prefixes that end inside L contribute L.cnt[a][b]. Prefixes that end inside R first pass through all of L, arriving at remainder L.g[a], then behave like R entered with that remainder — contributing R.cnt[L.g[a]][b]. Hence:

cnt[a][b] = L.cnt[a][b] + R.cnt[L.g[a]][b]
g[a]      = R.g[L.g[a]]

Store one such (cnt, g) at every node of a segment tree over the array. Each merge is O(k^3) (really O(k^2) counts plus the g composition), and k <= 5 keeps that a tiny constant.

Answering a query

1. Apply the persistent update by writing the leaf for index and re-merging up the tree — O(log n) merges.
2. Merge the nodes covering the range [start, n-1] into one aggregate block.
3. The empty product is 1, so we enter that block with remainder 1 mod k. The answer is cnt[1 mod k][x].

Worked example

nums = [1,2,3,4,5], k = 3, queries = [[2,2,0,2],[3,3,3,0],[0,1,0,1]].

• Query 0 [2,2,0,2]: set nums[2]=2 -> [1,2,2,4,5], start=0, x=2. Running products mod 3 of the prefixes are 1, 2, 1, 1, 2. Two of them equal 2, so answer 2.
• Query 1 [3,3,3,0]: set nums[3]=3 -> [1,2,2,3,5], start=3, x=0. Over [3,5] the prefix products mod 3 are 0, 0. Both equal 0, answer 2.
• Query 2 [0,1,0,1]: set nums[0]=1 (unchanged) -> [1,2,2,3,5], start=0, x=1. Prefix products mod 3 are 1, 2, 1, 0, 0. Two equal 1, answer 2.

Result: [2, 2, 2].

Complexity

• Time: O(n * k^2) to build, and O(log n * k^2) per update and per query, for a total of O((n + q log n) * k^2). With k <= 5 the k^2 factor is a small constant.
• Space: O(n * k^2) for the segment tree (each node stores a k x k count matrix and a length-k exit vector).

class Solution {
    int k, size;
    int[][][] cnt; // cnt[node][a][b]
    int[][] g;     // g[node][a]

    public int[] resultArray(int[] nums, int k, int[][] queries) {
        this.k = k;
        int n = nums.length;
        size = 1;
        while (size < n) size <<= 1;
        cnt = new int[2 * size][k][k];
        g = new int[2 * size][k];

        for (int i = 0; i < size; i++) {
            int m = (i < n) ? (nums[i] % k) : (1 % k);
            setLeaf(size + i, m);
        }
        for (int i = size - 1; i >= 1; i--) {
            merge(i, 2 * i, 2 * i + 1);
        }

        int startRem = 1 % k;
        int[] res = new int[queries.length];
        for (int qi = 0; qi < queries.length; qi++) {
            int index = queries[qi][0];
            int value = queries[qi][1];
            int start = queries[qi][2];
            int x = queries[qi][3];
            update(index, value % k);
            int[][] rc = queryRange(start, n - 1);
            res[qi] = rc[startRem][x];
        }
        return res;
    }

    private void setLeaf(int node, int m) {
        for (int a = 0; a < k; a++) {
            for (int b = 0; b < k; b++) cnt[node][a][b] = 0;
            int b = (a * m) % k;
            cnt[node][a][b] = 1;
            g[node][a] = b;
        }
    }

    // write merge of children L, R into node `dst`
    private void merge(int dst, int L, int R) {
        for (int a = 0; a < k; a++) {
            int la = g[L][a];
            for (int b = 0; b < k; b++) {
                cnt[dst][a][b] = cnt[L][a][b] + cnt[R][la][b];
            }
            g[dst][a] = g[R][la];
        }
    }

    private void update(int i, int m) {
        int p = size + i;
        setLeaf(p, m);
        p >>= 1;
        while (p >= 1) {
            merge(p, 2 * p, 2 * p + 1);
            p >>= 1;
        }
    }

    // merge nodes covering [l, r] inclusive; identity enters with g[a] = a
    private int[][] queryRange(int l, int r) {
        int[][] leftC = new int[k][k];
        int[] leftG = new int[k];
        for (int a = 0; a < k; a++) leftG[a] = a; // identity cnt is all zero

        // collect right-side nodes, apply in reverse
        int[] rightNodes = new int[2 * 32];
        int rn = 0;

        int lo = l + size;
        int hi = r + size + 1;
        while (lo < hi) {
            if ((lo & 1) == 1) {
                mergeInto(leftC, leftG, lo);
                lo++;
            }
            if ((hi & 1) == 1) {
                hi--;
                rightNodes[rn++] = hi;
            }
            lo >>= 1;
            hi >>= 1;
        }
        for (int idx = rn - 1; idx >= 0; idx--) {
            mergeInto(leftC, leftG, rightNodes[idx]);
        }
        return leftC;
    }

    // leftC/leftG := merge(left, tree[node])
    private void mergeInto(int[][] leftC, int[] leftG, int node) {
        int[][] newC = new int[k][k];
        int[] newG = new int[k];
        for (int a = 0; a < k; a++) {
            int la = leftG[a];
            for (int b = 0; b < k; b++) {
                newC[a][b] = leftC[a][b] + cnt[node][la][b];
            }
            newG[a] = g[node][la];
        }
        for (int a = 0; a < k; a++) {
            System.arraycopy(newC[a], 0, leftC[a], 0, k);
            leftG[a] = newG[a];
        }
    }
}