素数和组合问题の解法

問題概要

与えられたn個の整数と整数k(k < n)があり、その中からk個の数を選んで和をとったときに、その和が素数になる組み合わせ数を求める問題。

例としてn=4, k=3で、数列が3,7,12,19のとき、選べる組み合わせは4通りある:

  • 3+7+12 = 22
  • 3+7+19 = 29
  • 7+12+19 = 38
  • 3+12+19 = 34

このうち素数は29のみなので、出力は1となる。

入出力形式

入力形式:

4 3
3 7 12 19

出力形式:

1

解法の考え方

この問題は再帰やビット操作を使ってすべての組み合わせを試すことで解くことができる。

組み合わせの生成において重複を避けるため、「昇順原則(不降則)」を採用する。これにより、同じ要素の組み合わせが複数回カウントされるのを防ぐ。

アルゴリズム1:深さ優先探索(DFS)

DFSによる再帰的な探索方法。

  • m:選んだ要素数
  • sum:現在の和
  • s:次に選ぶ要素の開始位置

この方法では、選ぶ要素を常に昇順に進めることで、重複のない組み合わせを生成する。

#include <iostream>
using namespace std;

int n, k, a[25];
long long ans;

bool is_prime(int num) {
    if (num < 2) return false;
    for (int i = 2; i * i <= num; ++i) {
        if (num % i == 0) return false;
    }
    return true;
}

void dfs(int count, int total, int index) {
    if (count == k) {
        if (is_prime(total)) ans++;
        return;
    }
    for (int i = index; i < n; ++i) {
        dfs(count + 1, total + a[i], i + 1);
    }
}

int main() {
    cin >> n >> k;
    for (int i = 0; i < n; ++i) cin >> a[i];
    dfs(0, 0, 0);
    cout << ans;
    return 0;
}

アルゴリズム2:ビットマスクによる組み合わせ列挙

nビットのマスクを使って、選ぶ要素をビットで表現する。

  • __builtin_popcount()関数でビット中の1の数を数える。
  • 1のビット位置が選ぶ要素のインデックスに対応。
#include <iostream>
using namespace std;

int a[30];

bool is_prime(int num) {
    if (num < 2) return false;
    for (int i = 2; i * i <= num; ++i) {
        if (num % i == 0) return false;
    }
    return true;
}

int main() {
    int n, k, count = 0;
    cin >> n >> k;
    for (int i = 0; i < n; ++i) cin >> a[i];

    int max_mask = 1 << n;
    for (int mask = 0; mask < max_mask; ++mask) {
        if (__builtin_popcount(mask) == k) {
            int sum = 0;
            for (int i = 0; i < n; ++i) {
                if (mask & (1 << i)) sum += a[i];
            }
            if (is_prime(sum)) count++;
        }
    }

    cout << count;
    return 0;
}

タグ: C++ アルゴリズム 再帰 ビット操作 素数判定

5月15日 08:45 投稿