「mod 7 占い」(ランク S) のアルゴリズムを考えてみた

はじめに

paiza x Qiita のプログラミング問題 の mod 7 占い をトライしました. 複数のアルゴリズムで, 実行時間の比較をしてみました.

Google Colaboratory で実装したコード (の断片) は，末尾のリンクからたどれます.

素直なアルゴリズム

真面目に計算する

読み込んだ数列から 3 つを選ぶすべての組み合わせに対し, その和が 7 で割り切れるかどうかを計算します. コードはこんな感じ.

1-1

## データを読み込む
n_num = int(input().rstrip()) # 入力する数の個数
input_num = [] # 入力する数のリスト
for i in range(n_num): # 入力する数のリストを作成
    input_num.append(int(input().rstrip()))

# 数える
count = 0 # 答えを格納する変数

for i in range(n_num-2):
    for j in range(i+1, n_num-1):
        for k in range(j+1, n_num):
            sum = input_num[i] + input_num[j] + input_num[k]
            if sum % 7 == 0: # 7 で割った余りが 0 ならば
                count += 1   # インクリメント

print(count)

落とし穴があるかと思いましたが, これで正答率 100% でした.

計算の省力化

上記のコードの実行結果を見ると, 実行時間の長い場合がありました (平均 0.39s. 実行時間については「さいごに」にまとめました). そこで実行時間の短縮を検討します.

まず 3 つの数の和の計算についてです. 上記のコードでは毎回 3 つの数の和を計算しています. ここの部分について, まず最初の 2 つの数の和を計算しておき, その和に対して 3 番目の数を加えるように変更します.

変更後のコードは以下の通りです.

1-2

count = 0 # 答えを格納する変数

for i in range(n_num-2):
    for j in range(i+1, n_num-1):
        sum_ij = input_num[i] + input_num[j] # 最初の 2 つの数の和を計算しておく
        for k in range(j+1, n_num):
            sum =  sum_ij + input_num[k]     # 2 つの数の和に, 3 番目の数を加える
            if sum % 7 == 0:
                count += 1

print(count)

実行時間の平均は 0.30s となり, 少し短縮されました.

データ読み込み時の工夫

入力される数字の条件は $0 \le a_i \lt 2^{32}$ なので, 極端に大きな数が存在し, それを繰り返し計算すると時間がかかります. そこで入力された数字をそのまま使うのではなく, 7 で割った余りを計算しておいて, それを用いることで, 計算時間の短縮が図れないかと考えました.

データを読み込むところで, 7 で割った余りを保存します. コードは次の通りです.

1-3

n_num = int(input().rstrip()) # 入力する数の個数
input_num = [] # 入力する数のリスト
for i in range(n_num): # 入力する数のリストを作成
    input_num.append(int(input().rstrip()) % 7) # 7 で割った余りを保存する

データの入力部を上記 (1-3) のようにし, 計算部分を省力版 (1-2) にしたときの実行時間の平均は 0.23s でした. だいぶ短くなりました.

別のアプローチ

さらに時間短縮できないかと思い, 「組み合わせの数」を用いる方法を考えました.

例えば, 7 で割った余りが 1, 2, 4 となる 3 つの数の和は, 7 で割り切れます. ここで, 仮に入力された数のうち

• 余り 1 の数 ; 2 個
• 余り 2 の数 ; 3 個
• 余り 4 の数 ; 1 個

であれば, 余りが 1, 2, 4 になる 3 つの数の組み合わせの数は $2 \times 3 \times 1 = 6$ となります.

このように

1. 入力された数の余りを計算し, 余りが 0, 1, ..., 6 の数がそれぞれ何個あるか求める
2. 7 で割り切れる余りの組み合わせを計算する
3. 上記の組み合わせの数がそれぞれ何個あるかを計算し, 合算する

という考え方でも目的を達成することができます.

注意点としては, 例えば余りが 1, 1, 5 のように, 同じ余りが重複する場合です.

• 余り 1 の数 ; 4 個
• 余り 5 の数 ; 3 個

であれば, 余りが 1, 1, 5 となる組み合わせの数は ${}_4C_2 \times 3 = 18$ となります. 同じ余りが重複するかしないかで場合分けしながら計算します.

入力

まずデータの入力部です. num_list というリストを用意し, 余りが 0, 1, ..., 6 になる数字の個数を数えていきます.

2-1

num_list = [0] * 7 # 0で初期化

n_num = int(input().rstrip()) # 入力する数の個数
for i in range(n_num): # 7 で割った余りごとに, リストの要素を +1 する
    num_list[int(input().rstrip()) % 7] += 1

7 で割り切れる組み合わせ

次に, 7 で割り切れる余りの組み合わせを求めます. mod7_list というリストに 3 つの数のリストを保存していきます.

2-2

mod7_list = [] # 余りの和が 7 で割り切れる組み合わせを格納するリスト
for i in range(7):
    for j in range(i, 7):
        for k in range(j, 7):
            if (i + j + k) % 7 == 0:
                mod7_list.append([i, j, k])

ちなみに結果は次の通りになります.

2-2'

print(mod7_list)

[[0, 0, 0], [0, 1, 6], [0, 2, 5], [0, 3, 4], [1, 1, 5], [1, 2, 4], [1, 3, 3], [2, 2, 3], [2, 6, 6], [3, 5, 6], [4, 4, 6], [4, 5, 5]]

要注意は [0, 0, 0] の取り扱いですね.

解をまとめる

組み合わせの数の計算には, scipy.special.comb を使いました.

組み合わせの数を計算する際の場合分けは, 次のようにしています.

1. 該当する余りの数が 0 の場合 ; 組み合わせの数は 0 なので continue (もしかして, ここは無くても良いのかも)
2. 3 つの余りが同じとき (今回は 0, 0, 0 だけ) ; その余りの個数から 3 つを選ぶ計算をする ($={}_nC_3$)
3. 2 つの余りが同じとき (例えば 1, 1, 5) ; 同じ余りは組み合わせの数 ($={}_nC_2$) を計算し, そこに残りの余りの数を掛ける
4. 3 つの余りがすべて異なるとき ; 3 つの個数を掛け算する

コードは次の通りです.

2-3

from scipy.special import comb

count = 0 # 答えを格納する変数

for mod7 in mod7_list:
#    print(mod7)
    if num_list[mod7[0]] == 0 or num_list[mod7[1]] == 0 or num_list[mod7[2]] == 0:
        continue # 該当する余りの数が 0 の場合は continue

    if mod7[0] == mod7[1] and mod7[0] == mod7[2]: # 3 つの余りが同じとき
        sum = comb(num_list[mod7[0]], 3, exact=True)
    elif mod7[0] == mod7[1]:                      # 2 つの余りが同じとき (1)
        sum = comb(num_list[mod7[0]], 2, exact=True) * num_list[mod7[2]]
    elif mod7[1] == mod7[2]:                      # 2 つの余りが同じとき (2)
        sum = num_list[mod7[0]] * comb(num_list[mod7[1]], 2, exact=True)
    else:                                         # すべての余りが異なるとき
        sum = num_list[mod7[0]] * num_list[mod7[1]] * num_list[mod7[2]]

#    print(sum)
    count += sum

print(count)

一連のコードを実行した際の時間は, 平均 0.24s でした.

さいごに

各コードの実行時間をまとめます.

テスト No	1-1	1-2	1-3	2-1, 2, 3
1	0.03	0.03	0.02	0.26
2	0.03	0.03	0.03	0.22
3	0.03	0.03	0.03	0.23
4	0.03	0.03	0.04	0.25
5	0.07	0.06	0.06	0.22
6	0.37	0.27	0.24	0.23
7	1.39	1.05	0.84	0.22
8	0.08	0.07	0.06	0.24
9	0.43	0.33	0.24	0.24
10	1.43	1.11	0.77	0.25
平均	0.39	0.30	0.23	0.24
標本標準偏差	0.56	0.42	0.31	0.01

和を求めて 7 の余りを計算する方法 (1-*) はテスト No. によっては時間がかかっています. 一方, 和の余りが 7 になる組み合わせから計算する方法 (2-*) は標準偏差が小さくなってます. 入力データを確認していないのでここからは想像になりますが, データ数が多い場合でも実行時間がそれほど長くならないアルゴリズムになっていると思います.

参考

Google Colaboratory へのリンク.

-mod7_占い_1.ipynb
-mod7_占い_2.ipynb