包含不同素数的，长度最大为K的子序列的计数

问题描述

给定一个由正整数组成的序列 nums，和一个正整数 K，计算出所有包含不同素数且长度最大为 K 的子序列的个数。要求时间复杂度优于 $O(2^N)$，其中 $N$ 为序列的长度。

解决思路

1. 首先需要对于每个数求出其所有的素因数，可以使用试除法或线性筛法。
2. 然后考虑如何统计包含不同素数的子序列。使用状态压缩，对于每个数 $i$，使用二进制数表示其所有的素因数，例如 $10=2\times5$，则其二进制表示即为 $101$，$21=3\times7$，则其二进制表示为 $11000$，用 $state_i$ 表示第 $i$ 个数的状态，即所有素因数的二进制表示进行按位或运算得到。
3. 使用 DP 解决问题，设 $f_{i,s}$ 表示以第 $i$ 个数结尾，状态为 $s$ 的子序列个数。那么可以转移得到 $f_{i,s}=f_{i-1,s}+\sum f_{j, s;\text{xor};state_i}$，其中 $j<i$，表示以第 $i$ 个数结尾，且包含第 $i$ 个数的子序列个数。这里的 $\text{xor}$ 表示按位异或。

代码实现

def get_prime_factors(num):
    """返回 num 的素因数"""
    factors = []
    for i in range(2, int(num**0.5)+1):
        if num % i == 0:
            while num % i == 0:
                num //= i
            factors.append(i)
    if num > 1:
        factors.append(num)
    return factors

def count_distinct_prime_subsequences(nums, K):
    # 计算状态
    states = [0] * len(nums)
    for i, num in enumerate(nums):
        factors = get_prime_factors(num)
        state = 0
        for factor in factors:
            state |= 1 << (factor-2)
        states[i] = state

    # DP
    f = [[0] * (1 << (K-1)) for _ in range(len(nums))]
    for i in range(len(nums)):
        f[i][states[i]] = 1
        for s in range(1 << (K-1)):
            f[i][s] += f[i-1][s]
            for j in range(i):
                s2 = s ^ states[i]
                f[i][s] += f[j][s2]

    # 统计答案
    ans = 0
    for i in range(len(nums)):
        for s in range(1 << (K-1)):
            if bin(s).count('1') == K:
                ans += f[i][s]
    return ans

时间复杂度分析

• 对于每个数求素因数需要的时间复杂度为 $O(\sqrt{N}\log N)$，其中 $N$ 是数的大小的最大值。
• 状态压缩和 DP 的时间复杂度都是 $O(N\times 2^{K-1})$。
• 统计答案的时间复杂度是 $O(N\times 2^{K-1})$。
• 因此总时间复杂度是 $O(N\times 2^{K-1}+\sqrt{N}\log N)$，优于 $O(2^N)$。