かっつのメモ帳

主に競プロ 時々日記

April Lunchtime 2021 - IMAT

問題のリンク

問題概要

$N$行$ M $列の行列 $A$ が与えられる。この行列の部分行列であって、その総和が $B$ に等しくなるものの総数を $10^{9}+7$ で割った余りを求めよ。

部分行列とは、ある行集合と列集合があって、その行、及び列を取り除くことで得られる行列とする。

制約

  • $ 1 \leq N \cdot M \leq 200$

  • $1 \leq A[i][j],B \leq 10^{5}$

解法

転置行列を取るor取らないの前処理により、予め $ N \leq M $ にしておく。

$ 15^{2} \gt 200$ より、この時 $1 \leq N \leq 14$ である。

部分行列に含まれる行集合を全て試すことで、1次元の部分和問題の話に帰着させることが出来る。

毎回総和と、どこまで見たかの状態を持ったDPを解くことで、 $O(2^{N} MB)$

これで部分点2まで取れるが、満点制約の元ではTLEになってしまう。

ここで場合分けを行う。

  • $ M $ が小さい時

DPの代わりに半分全列挙をすることで、$O(2^{N} 2^{M/2} log 2^{M/2})$

数え上げるパートは尺取り法の要領で見ていけば良い。

$2^{14} * 2^{20/2} = 16777216$ から見積もって、$M \leq 20$ 程度では十分高速に動くことが予想される。

追記:tuteさんに計算量のlogを落とす方法を教えて貰いました。$ B $ 以下の値を何個作れるかだけが分かれば良いので、長さ $B+1$ の配列を使い回せば良い。

  • $ M $ が大きい時

$M \gt 20$ の場合を想定する。この時 $N \lt 10$ である。

この場合は先に挙げたDPをすれば良い。

$2^{9} * (200//9) * 10^{5} \approx 10^{9}$ とやや最悪ケースでは重そうだが、このDPは枝狩りが効きやすい上に定数倍的にも軽い処理なので多分大丈夫。


この2つの方針を組み合わせることで満点を得ることが出来る。

とはいえTLは結構厳しめ…(インド特有の攻めた0.50sec設定) c++17で出したら通った。

実装

#pragma GCC optimize("Ofast")
#include <bits/stdc++.h>
using namespace std;
typedef long long int ll;

constexpr ll mod=1e9+7;

void add(int &a,int b){
    a+=b;
    if(a>=mod)a-=mod;
}

int A[201][201];

int main(){
    cin.tie(nullptr);
    ios::sync_with_stdio(false);
    int q; cin >> q;
    while(q--){
        int n,m,B; cin >> n >> m >> B;
        for(int i=0;i<n;i++){
            for(int j=0;j<m;j++){
                cin >> A[i][j];
            }
        }
        if(n>m){
            for(int i=0;i<n;i++){
                for(int j=0;j<m;j++){
                    if(i>j)swap(A[i][j],A[j][i]);
                }
            }
            swap(n,m);
        }
        int res=0;
        vector<int> v(m);
        vector<int> dp(B+1);
        for(int i=1;i<(1<<n);i++){
            for(int j=0;j<m;j++){
                v[j]=0;
                for(int k=0;k<n;k++){
                    if((1<<k)&i){
                        v[j]+=A[k][j];
                    }
                }
            }
            if(m<=20){ // 半分全列挙
                int ret=0;
                vector<int> v1,v2;
                int m1=m/2,m2=m-m1;
                for(int j=0;j<(1<<m1);j++){
                    int sum=0;
                    for(int k=0;k<m1;k++){
                        if((1<<k)&j)sum+=v[k];
                    }
                    v1.push_back(sum);
                }
                for(int j=0;j<(1<<m2);j++){
                    int sum=0;
                    for(int k=0;k<m2;k++){
                        if((1<<k)&j)sum+=v[m1+k];
                    }
                    v2.push_back(sum);
                }
                sort(v1.begin(), v1.end());
                sort(v2.begin(), v2.end());
                int j=0,k=v2.size()-1;
                while(j<v1.size() and 0<=k){
                    int x=v1[j],cnt1=0;
                    while(j<v1.size() and v1[j]==x){
                        cnt1++; j++;
                    }
                    while(0<=k and v2[k]+x>B)k--;
                    if(0<=k and v2[k]+x==B){
                        int cnt2=0;
                        while(0<=k and v2[k]==B-x){
                            cnt2++; k--;
                        }
                        ret+=cnt1*cnt2;
                    }
                }
                add(res,ret);
            }
            else{
                fill(dp.begin(), dp.end(),0);
                dp[0]=1;
                for(int j=0;j<m;j++){
                    if(v[j]>B)continue;
                    for(int k=B;k-v[j]>=0;k--){
                        add(dp[k],dp[k-v[j]]);
                    }
                }
                add(res,dp[B]);
            }
        }
        cout << res << endl;
    }
}

感想

値の大小によって方針を変える問題は、AOJ-ICPCでも見たことがあったので思いつきたかった。