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

问题描述

给定一个由正整数构成的序列，找到长度不超过 K 的包含不同素数元素的子序列的数量。例如，如果序列为 {2,3,5,7,11}，K 的值为 3，则符合条件的子序列有 {2,3,5}，{2,3,7}，{2,3,11}，{2,5,7}，{2,5,11}，{2,7,11}，{3,5,7}，{3,5,11}，{3,7,11}，{5,7,11}，共 10 个。

解决方案

可以使用动态规划的思想来解决此问题。

令 $dp[i][j]$ 表示在前 $i$ 个元素中选择长度不超过 $j$ 的包含不同素数元素的子序列的数量。

转移方程如下：

$$ dp[i][j] = \begin{cases} dp[i-1][j] & 当前元素不是质数 \ dp[i-1][j-1]+dp[i-1][j-2]+\cdots+dp[i-1][0] & 当前元素是质数 \end{cases} $$

其中，第二种情况是因为考虑到当前元素选择时包括它和不包括它两种情况。

最终的答案就是 $dp[n][1]+dp[n][2]+\cdots+dp[n][k]$，其中 $n$ 是序列长度。

由于本题中需要大量计算素数，可以在预处理的时候把素数全部生成出来，减小计算量。

代码实现

以下是使用 Python3 语言实现的代码：

def count_subs(n, k, arr):
    def is_prime(num):
        if num == 2:
            return True
        elif num < 2 or num % 2 == 0:
            return False
        for i in range(3, int(num ** 0.5) + 1, 2):
            if num % i == 0:
                return False
        return True

    primes = [i for i in range(n+1) if is_prime(i)]

    dp = [[0] * (k+1) for _ in range(n+1)]
    for i in range(n+1):
        dp[i][0] = 1

    for i in range(1, n+1):
        for j in range(1, k+1):
            if not is_prime(arr[i-1]):
                dp[i][j] = dp[i-1][j]
            else:
                dp[i][j] = sum(dp[i-1][q] for q in range(j-1, -1, -1) if (q == 0 or primes.index(arr[i-1]) - primes.index(arr[i-q-1]) > q))

    return sum(dp[n][i] for i in range(1, k+1))

以上是算法实现部分的代码片段，完整代码请参见下方参考资料。

复杂度分析

算法的时间复杂度为 $O(nk^2)$，空间复杂度为 $O(nk)$。其中，$n$ 是序列长度，$k$ 是子序列最大长度。

参考资料

• Count of subsequence of length 1 to K whose elements are prime numbers