LeetCode in Kotlin

3901. Good Subsequence Queries

Hard

You are given an integer array nums of length n and an integer p.

A non-empty subsequence of nums is called good if:

You are also given a 2D integer array queries of length q, where each queries[i] = [indi, vali] indicates that you should update nums[indi] to vali.

After each query, determine whether there exists any good subsequence in the current array.

Return the number of queries for which a good subsequence exists.

The term gcd(a, b) denotes the greatest common divisor of a and b.

Example 1:

Input: nums = [4,8,12,16], p = 2, queries = [[0,3],[2,6]]

Output: 1

Explanation:

i [ind_i, val_i] Operation Updated nums Any good Subsequence
0 [0, 3] Update nums[0] to 3 [3, 8, 12, 16] No, as no subsequence has GCD exactly p = 2
1 [2, 6] Update nums[2] to 6 [3, 8, 6, 16] Yes, subsequence [8, 6] has GCD exactly p = 2

Thus, the answer is 1.

Example 2:

Input: nums = [4,5,7,8], p = 3, queries = [[0,6],[1,9],[2,3]]

Output: 2

Explanation:

i [ind_i, val_i] Operation Updated nums Any good Subsequence
0 [0, 6] Update nums[0] to 6 [6, 5, 7, 8] No, as no subsequence has GCD exactly p = 3
1 [1, 9] Update nums[1] to 9 [6, 9, 7, 8] Yes, subsequence [6, 9] has GCD exactly p = 3
2 [2, 3] Update nums[2] to 3 [6, 9, 3, 8] Yes, subsequence [6, 9, 3] has GCD exactly p = 3

Thus, the answer is 2.

Example 3:

Input: nums = [5,7,9], p = 2, queries = [[1,4],[2,8]]

Output: 0

Explanation:

i [ind_i, val_i] Operation Updated nums Any good Subsequence
0 [1, 4] Update nums[1] to 4 [5, 4, 9] No, as no subsequence has GCD exactly p = 2
1 [2, 8] Update nums[2] to 8 [5, 4, 8] No, as no subsequence has GCD exactly p = 2

Thus, the answer is 0.

Constraints:

Solution

class Solution {
    private lateinit var tree: IntArray
    private var validCount = 0

    fun countGoodSubseq(nums: IntArray, p: Int, queries: Array<IntArray>): Int {
        val n = nums.size
        // 1. Iterative Segment Tree only needs 2n space
        tree = IntArray(2 * n)
        // Build the leaves of the tree and get initial count
        for (i in 0..<n) {
            if (nums[i] % p == 0) {
                tree[n + i] = nums[i]
                validCount++
            }
        }
        // Build the internal nodes bottom-up
        for (i in n - 1 downTo 1) {
            // tree[i << 1] is the left child, tree[i << 1 | 1] is the right child
            tree[i] = gcd(tree[i shl 1], tree[i shl 1 or 1])
        }
        var ans = 0
        for (q in queries) {
            val idx = q[0]
            val value = q[1]
            // 2. O(1) tracking for validCount (No segment tree needed for this!)
            val wasValid = (nums[idx] % p == 0)
            val isValid = (value % p == 0)
            if (wasValid && !isValid) {
                validCount--
            }
            if (!wasValid && isValid) {
                validCount++
            }
            // Update the original array to keep track of the old values
            nums[idx] = value
            // Point update for the Iterative Tree
            tree[idx + n] = if (isValid) value else 0
            // Climb up the tree using bitwise shifts (i >>= 1 means i = i / 2)
            var i = idx + n
            while (i > 1) {
                tree[i shr 1] = gcd(tree[i], tree[i xor 1])
                i = i shr 1
            }
            // 3. The logic check (tree[1] is ALWAYS the root in an iterative tree)
            if (tree[1] == p) {
                if (validCount < n) {
                    ans++
                } else {
                    // validCount == n (Every element is a multiple of p)
                    if (n > 20) {
                        // THE O(1) MATH BYPASS!
                        ans++
                    } else {
                        // Only run this heavy check if n is 20 or smaller
                        var flag = false
                        for (i in 0..<n) {
                            val leftGcd = query(0, i - 1, n)
                            val rightGcd = query(i + 1, n - 1, n)
                            if (gcd(leftGcd, rightGcd) == p) {
                                flag = true
                                break
                            }
                        }
                        if (flag) {
                            ans++
                        }
                    }
                }
            }
        }
        return ans
    }

    // Iterative Range Query [l, r] inclusive
    private fun query(l: Int, r: Int, n: Int): Int {
        var l = l
        var r = r
        if (l > r) {
            return 0
        }
        var res = 0
        l += n
        r += n + 1
        while (l < r) {
            if ((l and 1) == 1) {
                res = gcd(res, tree[l++])
            }
            if ((r and 1) == 1) {
                res = gcd(res, tree[--r])
            }
            l = l shr 1
            r = r shr 1
        }
        return res
    }

    private fun gcd(a: Int, b: Int): Int {
        var a = a
        var b = b
        while (b > 0) {
            val temp = b
            b = a % b
            a = temp
        }
        return a
    }
}
This site is open source. Improve this page.