[백준]1415번 사탕 - C++

https://www.acmicpc.net/problem/1415

1415번: 사탕

문제

다솜이는 슈퍼에서 사탕을 사려고 한다. 슈퍼에는 사탕이 N종류가 있다. 각각의 사탕은 가격이 있다. 다솜이는 사탕을 사는데, 사탕의 가격의 합이 소수가 되게하려고 한다.

가격이 같은 사탕은 모양이 같게 생겼다. 따라서 다솜이는 사탕을 적절히 샀을 때, 그 모양이 전부 똑같은 방법은 사지 않으려고 한다.

예를 들어, (1, 2, 1, 3, 1)을 사는 것과, (3, 1, 1, 1, 2)를 사는 것은 같은 방법이다. 따라서 한번만 세야 한다.

입력

첫째 줄에 슈퍼에 있는 사탕의 개수 N이 주어진다. N은 50보다 작거나 같은 자연수이다. 둘째 줄부터 N개의 줄에 각 사탕의 가격이 주어진다. 사탕의 가격은 10,000보다 작거나 같은 음이 아닌 정수이다.

출력

첫째 줄에 다솜이가 사탕을 살 수 있는 방법의 경우의 수를 출력한다.

---

     dp 배열을 활용하여 문제를 풀 수 있다. 처음 입력을 받을 때 <가격, 개수>의 벡터로 입력을 받는다. 그 다음 냅색 유형의 풀이와 같이 dp[가격]을 활용한다. 벡터를 순회하며 최댓값 (n*10000)부터 dp[i] += dp[i - 가격 * 고를 개수] 를 한다. 이렇게 되면 ( i - 가격 * 고를 개수)에서 해당 사탕을 고를 개수만큼 고른 경우가 그대로 추가되므로 고를 개수를 바꿔가며 dp를 누적시켜 나가면 결국 dp표가 완성된다. 그 다음 dp 표를 순회하며 가격이 소수인 경우만 최종 결과에 더한다. 이때 사탕의 가격이 0 인 것은 최종 결과에 가격이 0 인 사탕을 개수 범위 내에서 고를 수 있으므로 최종 결과에 (가격이 0 인 사탕의 개수 + 1)을 곱해준다.

#include <iostream>
#include <vector>
#include <algorithm>
#include <cstring>
#include <cmath>
using namespace std;

vector<pair<int, int>> arr;
int n,a,zcnt=1;
bool pmat[500000];
long long dp[500001];
long long ans = 0;

void make_pmat() {
	memset(pmat, true, sizeof(pmat));
	pmat[0] = false; pmat[1] = false;
	int num = 10000 * n;
	for (int i = 2; i <= sqrt(num); i++) {
		if (pmat[i] == false) continue;
		for (int j = i * i; j <= num; j+=i) {
			pmat[j] = false;
		}
	}
}

int main() {
	cin >> n;
	for (int i = 0; i < n; i++) {
		int flag = -1;
		cin >> a;
		if (a == 0) { zcnt++; continue; }
		for (int i = 0; i < arr.size(); i++) {
			if (a == arr[i].first) {
				flag = i; break;
			}
		}
		if (flag == -1) arr.push_back({ a,1 });
		else arr[flag].second++;
	}
	make_pmat();

	dp[0] = 1;
	for (auto cur : arr) {
		for (int i = n * 10000; i >= 0; i--) {
			for (int j = 1; j <= cur.second; j++) {
				
				if (i - cur.first * j < 0) break;
				dp[i] += dp[i - cur.first * j];
				
			}
		}
	}

	for (int i = 2; i <= n * 10000; i++) {
		if (pmat[i]) ans += dp[i];
	}

	cout << ans*zcnt;
}

---

    처음에는 dfs를 활용하여 문제를 풀려 하였다. 코드 작성 후 예제 코드는 잘 나와 1가지 가격이 50개 있거나 50가지 가격이 1개 있거나 결국 a 가지 가격이 b 개 있을 때 a*b 가 최대 50이라고 생각해서 충분히 시간 안에 들어올 줄 알았지만 시간 초과가 나왔다.

void dfs(int k,int num) {
	if (k == arr.size()) {
		//cout << num << "\n";
		if (pmat[num]) ans++;
		return;
	}
	for (int i = 0; i <= arr[k].second; i++) {
		dfs(k + 1, num + i*arr[k].first);
	}

}

재귀 함수는 시간 복잡도 계산하기가 어려운 것 같다.