LeetCode in Kotlin

1994. The Number of Good Subsets

Hard

You are given an integer array nums. We call a subset of nums good if its product can be represented as a product of one or more distinct prime numbers.

Return the number of different good subsets in nums modulo 109 + 7.

A subset of nums is any array that can be obtained by deleting some (possibly none or all) elements from nums. Two subsets are different if and only if the chosen indices to delete are different.

Example 1:

Input: nums = [1,2,3,4]

Output: 6

Explanation: The good subsets are:

Example 2:

Input: nums = [4,2,3,15]

Output: 5

Explanation: The good subsets are:

Constraints:

Solution

@Suppress("NAME_SHADOWING")
class Solution {
    private fun add(a: Long, b: Long): Long {
        var a = a
        a += b
        return if (a < MOD) a else a - MOD
    }

    private fun mul(a: Long, b: Long): Long {
        var a = a
        a *= b
        return if (a < MOD) a else a % MOD
    }

    private fun pow(a: Long, b: Long): Long {
        // a %= MOD;
        // b%=(MOD-1);//if MOD is prime
        var a = a
        var b = b
        var res: Long = 1
        while (b > 0) {
            if (b and 1L == 1L) {
                res = mul(res, a)
            }
            a = mul(a, a)
            b = b shr 1
        }
        return add(res, 0)
    }

    fun numberOfGoodSubsets(nums: IntArray): Int {
        val primes = intArrayOf(2, 3, 5, 7, 11, 13, 17, 19, 23, 29)
        val mask = IntArray(31)
        val freq = IntArray(31)
        for (x in nums) {
            freq[x]++
        }
        for (i in 1..30) {
            for (j in primes.indices) {
                if (i % primes[j] == 0) {
                    if (i / primes[j] % primes[j] == 0) {
                        mask[i] = 0
                        break
                    }
                    mask[i] = mask[i] or pow(2, j.toLong()).toInt()
                }
            }
        }
        val dp = LongArray(1024)
        dp[0] = 1
        for (i in 1..30) {
            if (mask[i] != 0) {
                for (j in 0..1023) {
                    if (mask[i] and j == 0 && dp[j] > 0) {
                        dp[mask[i] or j] = add(dp[mask[i] or j], mul(dp[j], freq[i].toLong()))
                    }
                }
            }
        }
        var ans: Long = 0
        for (i in 1..1023) {
            ans = add(ans, dp[i])
        }
        ans = mul(ans, pow(2, freq[1].toLong()))
        ans = add(ans, 0)
        return ans.toInt()
    }

    companion object {
        private const val MOD = (1e9 + 7).toLong()
    }
}