[Atcoder]ABC251 B - At Most 3 (Judge ver.) （DFS）

传送门：B - At Most 3 (Judge ver.) (atcoder.jp)

题目大意

给定一个长度为N的数组（N <= 300），每次从中取出至多3个数字，问所有可能的且小于等于W的和(sum)有多少种。

思路

dfs，一共3层，每层可能的结果为a[1 ~ N]或0（不选）。

为了防止重复选择，用vis数组防止重复，还可以用k（循环开始位置）来剪枝。

第二个剪枝就是当sum > W时直接退出，这时候已经不可能是答案了。

用一个bitset数组来记录下那些数字已经被加到答案中了，避免重复添加。

代码

#include <bits/stdc++.h>
#define int long long
using namespace std;
const int maxn = 310;
int a[maxn];
bitset<maxn * 10000> add;//表示已经被加到答案中 
bitset<maxn> vis;//表示这个位置的数字已经被选了 

int N, W;

int dfs(int u, int k, int sum)
{
	int res = 0;
	
	if(sum > W)return 0; //剪枝，返回一个零元 
	if(u == 4) //当已经选完了 3 层 
	{
		//确保每个结果只被加一次 
		if(0 < sum && sum <= W && !add[sum])return add[sum] = 1; //这里会返回 ，同时会给add[sum]赋值 
		return 0;
	}
	res += dfs(u + 1, k, sum); //这一层不选 
	for(int i = k;i <= N; ++ i) //从k开始循环，可以剪枝减少重复遍历 
	{
		if(!vis[i]) //只能选择未被选过的 
		{
			vis[i] = 1;
			res += dfs(u + 1, i + 1, sum + a[i]);
			vis[i] = 0;
		}
	}
	return res;
}

signed main()
{
	cin >> N >> W;
	for(int i = 1;i <= N; ++ i)scanf("%lld", a + i);
	cout << dfs(1, 1, 0) << '\n';
	return 0;
}