3525. Find X Value of Array II

import java.util.Arrays;

class Solution {
    // Segment Tree Node definition
    static class Node {
        long prod;
        long[] remain;

        Node(int k) {
            this.prod = 1;
            this.remain = new long[k];
        }
    }

    private int n;
    private int K;
    private Node[] tree;

    public int[] resultArray(int[] nums, int k, int[][] queries) {
        this.n = nums.length;
        this.K = k;
        this.tree = new Node[4 * n];
        
        // Build the initial segment tree
        build(nums, 0, 0, n - 1);

        int[] ans = new int[queries.length];

        for (int i = 0; i < queries.length; i++) {
            int index = queries[i][0];
            int value = queries[i][1] % K;
            int start = queries[i][2];
            int x = queries[i][3];

            // Perform persistent point update
            update(0, 0, n - 1, index, value);

            // Query the suffix range from 'start' to the end of the array
            Node resNode = query(0, 0, n - 1, start, n - 1);
            ans[i] = (int) resNode.remain[x];
        }

        return ans;
    }

    private Node merge(Node left, Node right) {
        Node parent = new Node(K);
        parent.prod = (left.prod * right.prod) % K;

        // 1. Prefixes that stay completely in the left child
        for (int r = 0; r < K; r++) {
            parent.remain[r] = left.remain[r];
        }

        // 2. Prefixes that span through the left child and extend into the right child
        for (int r = 0; r < K; r++) {
            int combinedRemainder = (int) ((left.prod * r) % K);
            parent.remain[combinedRemainder] += right.remain[r];
        }

        return parent;
    }

    private void build(int[] nums, int cur, int left, int right) {
        tree[cur] = new Node(K);
        if (left == right) {
            int val = nums[left] % K;
            tree[cur].prod = val;
            tree[cur].remain[val] = 1;
            return;
        }
        int mid = left + (right - left) / 2;
        build(nums, 2 * cur + 1, left, mid);
        build(nums, 2 * cur + 2, mid + 1, right);
        tree[cur] = merge(tree[2 * cur + 1], tree[2 * cur + 2]);
    }

    private void update(int cur, int lo, int hi, int idx, int val) {
        if (lo == hi) {
            Arrays.fill(tree[cur].remain, 0);
            tree[cur].prod = val;
            tree[cur].remain[val] = 1;
            return;
        }
        int mid = lo + (hi - lo) / 2;
        if (idx <= mid) {
            update(2 * cur + 1, lo, mid, idx, val);
        } else {
            update(2 * cur + 2, mid + 1, hi, idx, val);
        }
        tree[cur] = merge(tree[2 * cur + 1], tree[2 * cur + 2]);
    }

    private Node query(int cur, int lo, int hi, int ql, int qr) {
        if (ql <= lo && hi <= qr) {
            return tree[cur];
        }
        int mid = lo + (hi - lo) / 2;
        if (qr <= mid) {
            return query(2 * cur + 1, lo, mid, ql, qr);
        }
        if (ql > mid) {
            return query(2 * cur + 2, mid + 1, hi, ql, qr);
        }
        Node leftNode = query(2 * cur + 1, lo, mid, ql, qr);
        Node rightNode = query(2 * cur + 2, mid + 1, hi, ql, qr);
        return merge(leftNode, rightNode);
    }
}

class Solution:
    class Node:
        def __init__(self, k):
            self.prod = 1
            self.remain = [0] * k

    def resultArray(self, nums, k, queries):
        n = len(nums)
        self.k = k
        self.tree = [None] * (4 * n)

        def merge(left, right):
            parent = self.Node(k)
            parent.prod = (left.prod * right.prod) % k

            for r in range(k):
                parent.remain[r] = left.remain[r]

            for r in range(k):
                combined = (left.prod * r) % k
                parent.remain[combined] += right.remain[r]

            return parent

        def build(node, lo, hi):
            self.tree[node] = self.Node(k)

            if lo == hi:
                val = nums[lo] % k
                self.tree[node].prod = val
                self.tree[node].remain[val] = 1
                return

            mid = (lo + hi) // 2
            build(2 * node + 1, lo, mid)
            build(2 * node + 2, mid + 1, hi)

            self.tree[node] = merge(
                self.tree[2 * node + 1],
                self.tree[2 * node + 2]
            )

        def update(node, lo, hi, idx, val):
            if lo == hi:
                self.tree[node].remain = [0] * k
                self.tree[node].prod = val
                self.tree[node].remain[val] = 1
                return

            mid = (lo + hi) // 2

            if idx <= mid:
                update(2 * node + 1, lo, mid, idx, val)
            else:
                update(2 * node + 2, mid + 1, hi, idx, val)

            self.tree[node] = merge(
                self.tree[2 * node + 1],
                self.tree[2 * node + 2]
            )

        def query(node, lo, hi, ql, qr):
            if ql <= lo and hi <= qr:
                return self.tree[node]

            mid = (lo + hi) // 2

            if qr <= mid:
                return query(2 * node + 1, lo, mid, ql, qr)

            if ql > mid:
                return query(2 * node + 2, mid + 1, hi, ql, qr)

            left = query(2 * node + 1, lo, mid, ql, qr)
            right = query(2 * node + 2, mid + 1, hi, ql, qr)

            return merge(left, right)

        build(0, 0, n - 1)

        ans = []

        for index, value, start, x in queries:
            update(0, 0, n - 1, index, value % k)
            res = query(0, 0, n - 1, start, n - 1)
            ans.append(res.remain[x])

        return ans

class Solution {
    struct Node {
        long long prod;
        vector<long long> remain;

        Node(int k) : prod(1), remain(k, 0) {}
    };

    int n, K;
    vector<Node*> tree;

    Node* merge(Node* left, Node* right) {
        Node* parent = new Node(K);

        parent->prod = (left->prod * right->prod) % K;

        for (int r = 0; r < K; r++) {
            parent->remain[r] = left->remain[r];
        }

        for (int r = 0; r < K; r++) {
            int combined = (left->prod * r) % K;
            parent->remain[combined] += right->remain[r];
        }

        return parent;
    }

    void build(vector<int>& nums, int node, int lo, int hi) {
        tree[node] = new Node(K);

        if (lo == hi) {
            int val = nums[lo] % K;
            tree[node]->prod = val;
            tree[node]->remain[val] = 1;
            return;
        }

        int mid = (lo + hi) / 2;

        build(nums, 2 * node + 1, lo, mid);
        build(nums, 2 * node + 2, mid + 1, hi);

        tree[node] = merge(
            tree[2 * node + 1],
            tree[2 * node + 2]
        );
    }

    void update(int node, int lo, int hi, int idx, int val) {
        if (lo == hi) {
            fill(tree[node]->remain.begin(),
                 tree[node]->remain.end(), 0);

            tree[node]->prod = val;
            tree[node]->remain[val] = 1;
            return;
        }

        int mid = (lo + hi) / 2;

        if (idx <= mid)
            update(2 * node + 1, lo, mid, idx, val);
        else
            update(2 * node + 2, mid + 1, hi, idx, val);

        tree[node] = merge(
            tree[2 * node + 1],
            tree[2 * node + 2]
        );
    }

    Node* query(int node, int lo, int hi, int ql, int qr) {
        if (ql <= lo && hi <= qr)
            return tree[node];

        int mid = (lo + hi) / 2;

        if (qr <= mid)
            return query(2 * node + 1, lo, mid, ql, qr);

        if (ql > mid)
            return query(2 * node + 2, mid + 1, hi, ql, qr);

        Node* left =
            query(2 * node + 1, lo, mid, ql, qr);

        Node* right =
            query(2 * node + 2, mid + 1, hi, ql, qr);

        return merge(left, right);
    }

public:
    vector<int> resultArray(
        vector<int>& nums,
        int k,
        vector<vector<int>>& queries
    ) {
        n = nums.size();
        K = k;

        tree.resize(4 * n);

        build(nums, 0, 0, n - 1);

        vector<int> ans;

        for (auto& q : queries) {
            int index = q[0];
            int value = q[1] % K;
            int start = q[2];
            int x = q[3];

            update(0, 0, n - 1, index, value);

            Node* res =
                query(0, 0, n - 1, start, n - 1);

            ans.push_back(res->remain[x]);
        }

        return ans;
    }
};

class Solution {
    constructor() {
        this.K = 0;
        this.tree = [];
    }

    merge(left, right) {
        const parent = {
            prod: (left.prod * right.prod) % this.K,
            remain: new Array(this.K).fill(0)
        };

        for (let r = 0; r < this.K; r++) {
            parent.remain[r] = left.remain[r];
        }

        for (let r = 0; r < this.K; r++) {
            const combined = (left.prod * r) % this.K;
            parent.remain[combined] += right.remain[r];
        }

        return parent;
    }

    resultArray(nums, k, queries) {
        const n = nums.length;
        this.K = k;
        this.tree = new Array(4 * n);

        const build = (node, lo, hi) => {
            this.tree[node] = {
                prod: 1,
                remain: new Array(k).fill(0)
            };

            if (lo === hi) {
                const val = nums[lo] % k;
                this.tree[node].prod = val;
                this.tree[node].remain[val] = 1;
                return;
            }

            const mid = Math.floor((lo + hi) / 2);

            build(2 * node + 1, lo, mid);
            build(2 * node + 2, mid + 1, hi);

            this.tree[node] = this.merge(
                this.tree[2 * node + 1],
                this.tree[2 * node + 2]
            );
        };

        const update = (node, lo, hi, idx, val) => {
            if (lo === hi) {
                this.tree[node].remain =
                    new Array(k).fill(0);

                this.tree[node].prod = val;
                this.tree[node].remain[val] = 1;
                return;
            }

            const mid = Math.floor((lo + hi) / 2);

            if (idx <= mid) {
                update(2 * node + 1, lo, mid, idx, val);
            } else {
                update(2 * node + 2, mid + 1, hi, idx, val);
            }

            this.tree[node] = this.merge(
                this.tree[2 * node + 1],
                this.tree[2 * node + 2]
            );
        };

        const query = (node, lo, hi, ql, qr) => {
            if (ql <= lo && hi <= qr) {
                return this.tree[node];
            }

            const mid = Math.floor((lo + hi) / 2);

            if (qr <= mid) {
                return query(
                    2 * node + 1,
                    lo,
                    mid,
                    ql,
                    qr
                );
            }

            if (ql > mid) {
                return query(
                    2 * node + 2,
                    mid + 1,
                    hi,
                    ql,
                    qr
                );
            }

            const left = query(
                2 * node + 1,
                lo,
                mid,
                ql,
                qr
            );

            const right = query(
                2 * node + 2,
                mid + 1,
                hi,
                ql,
                qr
            );

            return this.merge(left, right);
        };

        build(0, 0, n - 1);

        const ans = [];

        for (const [index, value, start, x] of queries) {
            update(0, 0, n - 1, index, value % k);

            const res = query(
                0,
                0,
                n - 1,
                start,
                n - 1
            );

            ans.push(res.remain[x]);
        }

        return ans;
    }
}