PaizaのSランク相当問題（レベルアップ問題集）の次の問題を解説します．
3つの解法を解説したいと思います．
- 解説1（愚直解）
- 解説2（Paizaの解説と同じ方針）
- 解説3（DP）
問題概要
数字のリストが与えられる．その中から異なる3つを選んだときの和が7で割り切れる数の組み合わせの個数を求めよ．
解法1（愚直解）
愚直に与えられた数字のリストから3つを選んで和のmodを計算します．
N = int(input())
A = [int(input()) for _ in range(N)]
ans = 0
for i in range(N):
for j in range(i+1,N):
for k in range(j+1,N):
if (A[i]+A[j]+A[k])%7 == 0:
ans += 1
print(ans)
解法1の解説
問題文で言われたことをそのまま3重ループを使って実装しています．
同じ数字は使えないので2つ目，3つ目のループの開始をそれぞれi+1,j+1としています．
実行結果（解法1）
スコア 100
| テスト | 実行時間(s) |
|---|---|
| 1 | 0.02 |
| 2 | 0.03 |
| 3 | 0.02 |
| 4 | 0.03 |
| 5 | 0.06 |
| 6 | 0.3 |
| 7 | 1.21 |
| 8 | 0.07 |
| 9 | 0.39 |
| 10 | 1.29 |
テストケース7とテストケース10でやや時間がかかっていることがわかります．
また手元の環境で，テストケースを作って実行してみましたが，待てど暮らせど結果は表示されませんでした（TLE）．
import random
N = 100000
with open('input.txt', 'w') as file:
file.write(str(N)+"\n")
for i in range(N):
file.write(str(random.randint(0,2**32))+"\n")
$ N=10^5 $の3重ループになっており，$ N^3 = 10^{15}$のループを実行しているため，実行時間が膨大になっています．
（Paizaではテストケースが弱いようでスコア100を取得できています）
以降の解法ではこの部分も改善します．
解法2（Paizaの解説と同じ方針）
N = int(input())
A = [int(input())%7 for _ in range(N)]#7で割っておく
CNT = [0]*7
for a in A:CNT[a] += 1
ans = 0
for i in range(7):
for j in range(i,7):
for k in range(j,7):
if (i+j+k)%7 == 0:
if i == j and j == k:#3つとも同じ
ans += CNT[i]*(CNT[i]-1)*(CNT[i]-2)//6
elif i == j:#2個が同じ
ans += CNT[i]*(CNT[i]-1)//2 * CNT[k]
elif i == k:#2個が同じ
ans += CNT[i]*(CNT[i]-1)//2 * CNT[j]
elif j == k:#2個が同じ
ans += CNT[j]*(CNT[j]-1)//2 * CNT[i]
else:#3つとも違う
ans += CNT[i]*CNT[j]*CNT[k]
print(ans)
解法2の解説
高校数学の知識を使います．
和をとってからmodを計算してもmodをとってから和を計算しても結果が変わらないという事実を利用します．
簡単な例
11+12+13 = 36 = 1 (mod 7)
一方
11 = 4 (mod 7)
12 = 5 (mod 7)
13 = 6 (mod 7)
なので
11+12+13 = 4+5+6 = 15 = 1 (mod 7)
この事実により，各数字を7で割ったあまりに注目すれば，大きな数を扱う必要がなくなります．また，7で割ったあまりが同じ数を区別する必要もありません．
| 入力 | mod7 |
|---|---|
| 10 | 3 |
| 4 | 4 |
| 14 | 0 |
実際，元の入力で$ 10+4+14=28=0(\mod 7)$と計算しても，変換して
$3+4+0 = 7 = 0 (\mod 7)$と計算してもどちらも同じ結果になることがわかります．
したがって，与えられたリストの数字それぞれのmod 7を計算して，0~6の数字に変換しておきます．
そして，変換した入力リストに0~6がそれぞれいくつあるかを数えます．
次の例を考えます．
| 入力 | mod 7 |
|---|---|
| 1 | 1 |
| 2 | 2 |
| 3 | 3 |
| 4 | 4 |
| 5 | 5 |
| 6 | 6 |
| 7 | 0 |
| 8 | 1 |
| 9 | 2 |
| 10 | 3 |
すると，0~6の個数は次の様になります．
| 個数 | |
|---|---|
| 0 | 1 |
| 1 | 2 |
| 2 | 2 |
| 3 | 2 |
| 4 | 1 |
| 5 | 1 |
| 6 | 1 |
この状態で3つ足して和が0（または7)になるものを数えます．
例えば0,3,4を1つずつ選べば，0+3+4=7=0となり，条件を満たしています．
また，0は1個，3が2個，4が1個あるので，0+3+4という数字の組み合わせの作り方は（元の入力の数字を使えば）$ 1\times 2 \times1 = 2$ 通りあることがわかります．
では次に，2+2+3 = 7 = 0について考えましょう．この組み合わせの場合は
2は2個の中から2個選ぶ1通り，3は2個の中から1個選ぶ2通りの選び方があります．したがって，$1 \times 2 = 2$通りの選び方があることがわかります．
では，ある入力に対して，0~6の個数が次のようになったとします．
すると，0~6の個数は次の様になります．
| 個数 | |
|---|---|
| 0 | $a_0$ |
| 1 | $a_1$ |
| 2 | $a_2$ |
| 3 | $a_3$ |
| 4 | $a_4$ |
| 5 | $a_5$ |
| 6 | $a_6$ |
いくつかのパターンの組み合わせの個数を計算してみます．
-
0+3+4（全部異なる）
0は$a_0$個の中から1個選ぶので$a_0$通り．3, 4の選び方も同様に$a_3$通り，$a_4$通りです．
したがって，このときは$a_0 \times a_3 \times a_4$通りの組み合わせがあることがわかります． -
2+2+3（2個が同じ）
2の選び方は，$a_2$の中から2個選びます．これは実は，$a_2 \times (a_2-1) / 2$通りの選び方があります．（高校数学でいう組み合わせ， $_n C_r$です(後に簡単な解説を書きました)．）
3の選び方は$a_3$通りです．
したがって，$a_2\times(a_2-1)\times a_3 /2$ がこの場合の組み合わせの個数です． -
0+0+0（3個すべて同じ）
0は$a_0$個の中から3個選ぶので，$a_0\times(a_0-1)\times(a_0-2)/6$通りの選び方があります．
このような形で0~6の個数さえ解ってしまえば計算で組み合わせの個数がわかります．
ちなみに3つの和がmod 7 で0になる数の組み合わせはの次の11通りです．（パターンが漏れていたので追記しました、コメントでのご指摘ありがとうございました。）
- 0+0+0
- 0+1+6
- 0+2+5
- 0+3+4
- 1+2+4
- 1+3+3
- 2+2+3
- 2+6+6
- 3+5+6
- 4+4+6
- 4+5+5
この11通りに対して丁寧に計算しても良いですが，0~6から（重複を許して）適当に3つとってきて，和がmod 7 で0になるものだけ，先ほどの計算により組み合わせを求めてあげるという形でもよいと思います．
$_n C_r$の簡易な説明
5種類のアルファベット（A,B,C,D,E）の中から異なる2種類を選ぶことを考えます．
- まず１つ目のアルファベットを選びます．A~Eから選ぶので5通りです．
- 次に，１つ目に選んだアルファベット以外の4種類から1つ選ぶので4通りです．
- 最初の5通りの選び方それぞれに，２つ目のアルファベットの選び方4通りが対応するので，$4 \times 5 = 20$通りとなります．
- ただし，この20通りには次のようなものが含まれています．最初にA,次にBを選んだ[A,B]というペア．最初にB，次にAを選んだ[B,A]というペアを両方とも数えています．これらはアルファベットの組み合わせという点では同じものなので，数えすぎです．20通りすべてに対して，順番をひっくり返したペアが存在するので，実際に組み合わせの数はペアを考慮して，ちょうど半分です．よって，20/2 = 10 となります．
- これは先ほどの書き方で言うと5*(5-1)/2という計算です．
では3種類を選ぶ場合はどうでしょうか．
5種類のアルファベット（A,B,C,D,E）の中から異なる3種類を選ぶことを考えます．
- やはりまずは，順番に3種類選んでいって，543 = 60通りの選び方があります．
- しかし，先ほどと同じ様に，[A,B,C],[A,C,B],[B,A,C],[B,C,A],[C,A,B],[C,A,B]の6個は同じアルファベットの組み合わせです．必ず6個あるので，$5\times 4\times3=60$を6で割ってあげればよいです．
- したがって，この場合は5*(5-1)*(5-2)/6が求める答えとなります．
実行結果（解法2）
スコア 100
| テスト | 実行時間(s) |
|---|---|
| 1 | 0.03 |
| 2 | 0.03 |
| 3 | 0.03 |
| 4 | 0.03 |
| 5 | 0.03 |
| 6 | 0.03 |
| 7 | 0.03 |
| 8 | 0.03 |
| 9 | 0.03 |
| 10 | 0.03 |
すべてのテストケースで実行結果が0.03秒になりました．
では，さきほど手元環境では結果が返ってこなかったN=10000のデータでも試してみます．
解答：23808804962077
実行時間 = 0.013939380645751953(sec)
（入力受取後の処理の時間を計測）
となり，十分高速に答えが求まりました．
解法3（DP）
DPを用いた解法です．
N = int(input())
A = [int(input())%7 for _ in range(N)]
mod = 7
# DP[i][j] i 個の要素を選んでその和を7で割った余りが j であるような組み合わせの数
dp = [[0]*7 for _ in range(4)]
dp[0][0] = 1
for a in A:
tmp = [row[:] for row in dp]
for i in range(1,4):
for j in range(7):
tmp[i][(j+a)%mod] += dp[i-1][j]
dp = tmp
print(dp[3][0])
解法3の解説
DPの基本的な問題に帰着できます（2次元ですが）．次のようなDPテーブルを考えます．
- DP[i][j]:$i$個の数字（$i$枚のカード）を選んで，和（mod 7)が$j$であるようなものの組み合わせの数
DPテーブルの遷移
入力として，次が与えられたとします．（N=5）
| 入力 | mod7 |
|---|---|
| 2 | 2 |
| 3 | 3 |
| 14 | 0 |
| 6 | 6 |
| 5 | 5 |
するとDPテーブルの遷移は次のようになります．
また，実装上の注意として，あらたな入力$a$に対して，dp[0],dp[1],dp[2],dp[3]を同時に更新する必要がありますが，i=0,1,2,3と更新していくと,例えばdp[2]を更新したあとで，dp[3]更新することになり，dp[3]が正しく更新されません．（入力$a$の影響を受けたdp[2]の値を使ってしまう）
これを回避するために，tmpという一時的にDPテーブルをコピーするためのテーブルを用意しています．
実行結果（解法3）
スコア 100
| テスト | 実行時間(s) |
|---|---|
| 1 | 0.03 |
| 2 | 0.02 |
| 3 | 0.04 |
| 4 | 0.02 |
| 5 | 0.03 |
| 6 | 0.03 |
| 7 | 0.03 |
| 8 | 0.03 |
| 9 | 0.03 |
| 10 | 0.03 |
となりました．
また，N=10000のデータに対する結果はこちらです（解法2で使用したデータと全く同じ）
実行時間 = 1.0229029655456543(sec)
（入力受取後の処理の時間を計測）
解法2と比較してやや時間がかかっていますが，十分高速に動作しています．
これは解法2では，高々7回の3重ループの処理ですが，解法3ではN=100000のループを実行しているためです．
終わりに
今後の学習に向けて
ネット上にあるプログラムを読んでも意味がわからないと言うことがあった場合に，最近ではChatGPTを利用しています．
例えば，ChatGPTに解法2のコードを与えて説明してもらったものがこちらです（GPT3.5を使用しました）．
すると，次の様な結果が返ってきました．
長いので折りたたみ（ChatGPTの回答）
このプログラムは、与えられた数列の中から3つの要素を選んだときに、その和が7で割り切れる組み合わせの数を求めています。以下にその処理の詳細を説明します。
プログラムの各部分の説明
-
入力の取得と前処理
N = int(input()) A = [int(input())%7 for _ in range(N)] # 7で割った余りを計算-
Nは数列の長さを表します。 -
Aは長さNのリストで、各要素を7で割った余りを格納しています。この操作は、数列の和が7で割り切れる条件を検証しやすくするためです。
-
-
余りごとのカウント
CNT = [0]*7 for a in A: CNT[a] += 1-
CNTは長さ7のリストで、CNT[i]は余りがiである要素の個数を表します。
-
-
組み合わせの計算
ans = 0 for i in range(7): for j in range(i,7): for k in range(j,7): if (i+j+k)%7 == 0: if i == j and j == k: ans += CNT[i]*(CNT[i]-1)*(CNT[i]-2)//6 elif i == j: ans += CNT[i]*(CNT[i]-1)//2 * CNT[k] elif i == k: ans += CNT[i]*(CNT[i]-1)//2 * CNT[j] elif j == k: ans += CNT[j]*(CNT[j]-1)//2 * CNT[i] else: ans += CNT[i]*CNT[j]*CNT[k] print(ans)- 三重のループで
i,j,kの組み合わせを全てチェックします。 -
(i+j+k)%7 == 0である場合に限り、以下のようにansを更新します。-
i,j,kがすべて同じ場合：CNT[i]から3つ選ぶ組み合わせの数を計算します。 -
i,j,kのうち2つが同じ場合：CNT[i]から2つ選び、残り1つをCNT[k]から選ぶ組み合わせの数を計算します。 -
i,j,kがすべて異なる場合：それぞれのCNTの要素を掛け合わせます。
-
- 三重のループで
具体例を用いた説明
例えば、入力データが次のような場合を考えます。
N = 5
A = [14, 3, 2, 6, 5] # 入力データ
この場合、各要素を7で割った余りを求めると A = [0, 3, 2, 6, 5] となります。これを元に CNT を計算すると、
CNT = [1, 0, 1, 1, 0, 1, 1] # 余り0, 2, 3, 5, 6 の要素がそれぞれ1個ずつある
この場合の組み合わせは以下の通りです。
-
i = 0,j = 3,k = 4の場合、(0+3+4)%7 = 7%7 = 0 なので、CNT[0]*CNT[3]*CNT[4]を計算します。 -
i = 2,j = 5,k = 6の場合、(2+5+6)%7 = 13%7 = 6 なので、この組み合わせは無視されます。
プログラムはこのようにすべての組み合わせをチェックして、条件に合う組み合わせの数を数え上げ、最終的に ans にその数を出力します。
問題文を与えたわけではありませんが，ここまで解説してくれます．不明な点があれば更に詳しい説明を求めてもよいでしょう．
ちなみに，DPテーブルの遷移も聞いてみたところ，少し誤りを含んでいましたが，かなり参考になる回答が返ってきました．
ChatGPT等をうまく使って勉強していければいいのかなと思っています．
ここまでお読みいただいてありがとうございました．
誤り等あればご指摘ください．
少しでも参考になったという方は，ぜひいいねをお願いします．