paizaプログラミング問題 mod7占い

問題ページ

問題の概要

• 入力

  入力は以下のフォーマットで与えられます。

  $N$
  $a_1$
  $a_2$
  $\vdots$
  $a_N$

  $N$ は与えられるカードの枚数を表します。

  $a_i(1 \leq i \leq N)$はi 枚目のカードに書かれた整数であり、改行区切りで $N$ 個与えられます。

• 出力
  $1 \leq i < j < k \leq N$ で $a_i+a_j+a_k$ が7で割り切れるものの組み合わせ

解法

動的計画法で解く。
$i$枚目までで1枚選んだときに和のあまりが$j$となるような組み合わせの数をそれぞれ考えると、dpテーブルは
$$
dp1[i][j] = dp1[i-1][j] + (a[i]\mod 7 == j) \ \ (0 \leq i \leq 6).
$$
2枚の場合は、
$$
dp2[i][j] = dp2[i-1][j] + dp1[i-1][j-a[i]] \ \ (0 \leq i \leq 6).
$$
3枚の場合は2枚のときと同じで
$$
dp3[i][j] = dp3[i-1][j] + dp2[i-1][j-a[i]] \ \ (0 \leq i \leq 6).
$$
最後に $dp3[N][0]$ を出力すればよい。

コード例

solve.cpp

#include <bits/stdc++.h>
using namespace std;
#define rrep(i, a, b) for (int i = (a); i < (b); ++i)
#define rep(i, n) rrep(i,0,n)
typedef long long ll;
int main(void){
    ios::sync_with_stdio(false);
    cin.tie(nullptr);

    // 入力を受け取る

    ll N;
    cin >> N;

    vector<ll> a(N);
    rep(i,N)cin >> a[i];

    vector<vector<ll>> dp1(N+1,vector<ll>(7,0));
    vector<vector<ll>> dp2(N+1,vector<ll>(7,0));
    vector<vector<ll>> dp3(N+1,vector<ll>(7,0));


    rep(i,N){
        rep(j,7){
            dp1[i+1][j] = dp1[i][j];
            if(a[i]%7==j)dp1[i+1][j]++;
        }
        rep(j,7){
            ll k = 7000000000+j-a[i];
            k%=7;
            dp2[i+1][j] = dp2[i][j] + dp1[i][k];
        }
        rep(j,7){
            ll k = 7000000000+j-a[i];
            k%=7;
            dp3[i+1][j] = dp3[i][j] + dp2[i][k];
        }
    }

    cout << dp3[N][0];
    return 0;
}