【Paiza Sランク問題】mod7占い 解説【Python】

問題文

Paizaラーニング「mod7占い」

あなたは今、「mod7占い」というサービスを始めようと考えています。

mod7占いとは、整数が書かれた複数のカードの中から3枚を選び、そこに書かれた整数の和が7で割り切れるかどうかで運勢を決めようというものです。 カードは必ず異なる3枚を選ぶ必要があり、全てのカードには全て異なる数字が書かれています。

占いというからには、7で割り切れる組み合わせはそれなりに少なくする必要があります。 そこで、適当な複数のカードに対して、カードに書かれた3つの整数を足した和が7で割り切れるような組合せがいくつあるかを計算するプログラムを作成してください。

入力される値

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

N
a_1
a_2
...
a_N

• N は与えられるカードの枚数を表します。
• a_i (1 ≦ i ≦ N) はi 枚目のカードに書かれた整数であり、改行区切りでN 個与えられます。
• 入力値最終行の末尾に改行が１つ入ります。
• 文字列は標準入力から渡されます。

期待する出力

• 組合せ数を一行に出力してください。

条件

すべてのテストケースにおいて、以下の条件をみたします。

1 ≦ N ≦ 100000
0 ≦ a_i < 2^32

解答コード例

import sys
import itertools

# 3枚のカードの和が7で割り切れるかを判定する関数
def mod7_divisible_combinations(card_numbers):
    count = 0  # 7で割り切れる組み合わせの数をカウントする変数
    
    # itertools.combinationsを使って3つのカードの全ての組み合わせを生成する
    for combination in itertools.combinations(card_numbers, 3):
    
        if sum(combination) % 7 == 0:  # 組み合わせの和が7で割り切れるか確認する
            count += 1  # 条件を満たす場合、カウントを増やす
            
    return count


# メイン関数
def main():
    input = sys.stdin.read
    
    # 標準入力からデータを読み取る
    data = input().strip().split()
    N = int(data[0])  # 最初の要素はカードの枚数
    card_numbers = list(map(int, data[1:]))  # 残りの要素はカードに書かれた整数
    
    # 7で割り切れる組み合わせの数を計算
    result = mod7_divisible_combinations(card_numbers)
    
    # 結果を出力
    print(result)


if __name__ == "__main__":
    main()

解説

基本的な手順は以下。

1. 与えられた整数リストから3つの整数を選び出す全ての組み合わせを生成する。
2. 各組み合わせの和を計算し、その和が7で割り切れるかどうかを判定する。
3. 7で割り切れる組み合わせの数をカウントする。

具体例

入力例として次のようなカードのリストを考える。

N = 5
カードの整数: 1, 2, 3, 4, 5

1. 組み合わせの生成

   • まず、与えられた整数リストから3枚のカードを選ぶ全ての組み合わせを生成する。
   • これは、itertools.combinationsを使用することで実現できる。この関数は、与えられたリストから指定された数の要素を含む全ての組み合わせを生成する。
   • 具体例では、次の組み合わせが生成される。

   (1, 2, 3)
   (1, 2, 4)
   (1, 2, 5)
   (1, 3, 4)
   (1, 3, 5)
   (1, 4, 5)
   (2, 3, 4)
   (2, 3, 5)
   (2, 4, 5)
   (3, 4, 5)
   

    

2. 組み合わせの生成

   • 各組み合わせについて、3つの整数の和を計算する。
   • 和が7で割り切れるかどうかを確認する。割り切れる場合はカウントを増やす。
   • 具体例では、各組み合わせの和を計算し、7で割った余りを確認している。

   (1, 2, 3) -> 和: 6 -> 6 % 7 = 6 -> 割り切れない
   (1, 2, 4) -> 和: 7 -> 7 % 7 = 0 -> 割り切れる -> カウント+1
   (1, 2, 5) -> 和: 8 -> 8 % 7 = 1 -> 割り切れない
   (1, 3, 4) -> 和: 8 -> 8 % 7 = 1 -> 割り切れない
   (1, 3, 5) -> 和: 9 -> 9 % 7 = 2 -> 割り切れない
   (1, 4, 5) -> 和: 10 -> 10 % 7 = 3 -> 割り切れない
   (2, 3, 4) -> 和: 9 -> 9 % 7 = 2 -> 割り切れない
   (2, 3, 5) -> 和: 10 -> 10 % 7 = 3 -> 割り切れない
   (2, 4, 5) -> 和: 11 -> 11 % 7 = 4 -> 割り切れない
   (3, 4, 5) -> 和: 12 -> 12 % 7 = 5 -> 割り切れない
   

    

3. 結果の出力

   • 最終的に、7で割り切れる組み合わせの数を出力する。
   • 具体例では、1つの組み合わせ (1, 2, 4) のみが7で割り切れるため、結果は1となる。

計算量

組み合わせの数はN個の中から3個を選ぶ方法の数であるため、計算量は$nC_3 = \frac{n(n-1)(n-2)}{6}$、すなわち$O(N^3)$となる。
$N = 100000$のテストケースではおそらくタイムオーバーになるが、Paizaのテストケースは控えめに作られているようで、上記のコードでもクリアが可能。

解答コード例（高速化）

数の余りを利用して計算を簡略化することで、計算量を$O(N)$に削減できる。

1. 各カードの数を7で割った余りをあらかじめ計算しておく。
2. 余りの組み合わせを調べて、和が7で割り切れるかどうかを確認する。

def mod7_divisible_combinations(card_numbers):
    remainder_count = [0] * 7  # 余りごとのカウント

    # 各数値の余りを計算してカウント
    for number in card_numbers:
        remainder_count[number % 7] += 1

    count = 0

    # 余りの組み合わせを調べて和が7で割り切れるかを確認
    # 例: (0, 0, 0), (0, 1, 6), (1, 2, 4), (1, 3, 3) など
    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 == k:
                        count += remainder_count[i] * (remainder_count[i] - 1) * (remainder_count[i] - 2) // 6
                    elif i == j:
                        count += remainder_count[i] * (remainder_count[i] - 1) // 2 * remainder_count[k]
                    elif j == k:
                        count += remainder_count[i] * remainder_count[j] * (remainder_count[j] - 1) // 2
                    else:
                        count += remainder_count[i] * remainder_count[j] * remainder_count[k]

    return count

# 入力とメイン関数は同じ