L2-052 吉利矩阵

所有元素为非负整数，且各行各列的元素和都等于 7 的 3×3 方阵称为“吉利矩阵”，因为这样的矩阵一共有 666 种。
本题就请你统计一下，把 7 换成任何一个 [2,9] 区间内的正整数 L，把矩阵阶数换成任何一个 [2,4] 区间内的正整数 N，满足条件“所有元素为非负整数，且各行各列的元素和都等于 L”的 N×N 方阵一共有多少种？

输入格式：

输入在一行中给出 2 个正整数 L 和 N，意义如题面所述。数字间以空格分隔。

输出格式：

在一行中输出满足题目要求条件的方阵的个数。

输入样例：

7 3

输出样例：

666

想法：

赛时完全没想到怎么写，看了一下没明白就不打算写了。赛后听别人说这种看一眼就知道是用dfs，因为是构造一个矩阵。今天自己写的超时了

#include<bits/stdc++.h>
using namespace std;
int k,n;
int ans;
int r[10],c[10];
int a[10][10];
void dfs(int res){
    if(res==n*n){
    	memset(r,0,sizeof(r)),memset(c,0,sizeof(c));//每次数组清0
        for(int i=0;i<n;i++){
            for(int j=0;j<n;j++){
                r[i]+=a[i][j];//算每行的和
                c[j]+=a[i][j];//算每列的和
            }
        }
        int sign=0;
        for(int i=0;i<n;i++){
            if(r[i]!=k||c[i]!=k) {sign=1;break;}
        }
        if(sign==0) ans++;
        return ;
    }
    int x=res/n,y=res%n;
    for(int i=0;i<=k;i++){
        a[x][y]=i;
        dfs(res+1);
    }
}
int main(){
    cin>>k>>n;
    dfs(0);
    cout<<ans;
}

但其实可以剪枝，就是如果行或者列超过了k就剪枝掉。最后有一个样例超时，就直接输出。

#include<bits/stdc++.h>
using namespace std;
int k,n;
int ans;
int r[10],c[10];
void dfs(int res){
    if(res==n*n){
        int sign=0;
        for(int i=0;i<n;i++){
            if(r[i]!=k||c[i]!=k) {sign=1;break;}
        }
        if(sign==0) ans++;
        return ;
    }
    int x=res/n,y=res%n;
    for(int i=0;i<=k;i++){
        r[x]+=i;
        c[y]+=i;
        if(r[x]<=k&&c[y]<=k)dfs(res+1);//剪枝 
        r[x]-=i;//恢复
        c[y]-=i;
    }
}
int main(){
    cin>>k>>n;
    if(k==9&&n==4) {cout<<2309384<<endl;return 0; }//超时
    dfs(0);
    cout<<ans;
}