Pythonで無限数列から生成したベクトルの内積を求めるプログラム

問題の概要

3つの整数 c、m、n が与えられたとします。まず、次の条件に従う無限数列を生成することを考えます。

• 第1項は 0
• 第2項は c
• 第3項以降は漸化式 k_i = (k_i-2 + k_i-1) mod m に従う

この数列を第 2n+1 項まで生成します。続いて、数列の第3項以降の値を使い、連続する2つの値をそれぞれx座標・y座標とみなして二次元ベクトルを作成し、合計 n 個のベクトルを生成します。

さらに別の集合 S を用意します。S の各要素はベクトル i とベクトル j の内積（スカラー積）であり、1 ≤ i, j ≤ n かつ i ≠ j を満たす組み合わせが対象となります。最終的な答えは、集合 S に現れる異なる剰余値の個数です。内積の値が非常に大きくなる場合は mod m を適用します。

入出力例

たとえば、入力が 5, 6, 4 の場合、出力は 3 になります。

• 生成される数列：[0, 5, 5, 4, 3, 1, 4, 5, 3, 2]
• 生成されるベクトル：(5, 4)、(3, 1)、(4, 5)、(3, 2)
• 各ベクトル同士の内積を mod 6 で計算すると、集合 S には3種類の剰余値しか現れません。
• したがって、結果は 3 mod 6 = 3 となります。

解法のアルゴリズム

この問題を解くための手順は以下の通りです。

• n が 1 と等しい場合は 0 を返す
• それ以外の場合：
  • サイズ 2n+2 のリスト temp_arr を 0 で初期化する
  • temp_arr[0] := 0、temp_arr[1] := c と設定する
  • 空のリスト arr2 を用意する
  • i を 2 から 2n+1 まで動かしながら、temp_arr[i] := (temp_arr[i-1] + temp_arr[i-2]) mod m を計算する
  • i を 2 から 2n-3 まで 2 ずつ増やしながら：
    • temp := (temp_arr[i] × temp_arr[i+2] + temp_arr[i+1] × temp_arr[i+3]) mod m を計算し、arr2 の末尾に追加する
    • temp := (temp_arr[i] × temp_arr[i+4] + temp_arr[i+1] × temp_arr[i+5]) mod m を計算し、arr2 の末尾に追加する
  • temp := (temp_arr[2n-2] × temp_arr[2n] + temp_arr[2n-1] × temp_arr[2n+1]) mod m を計算し、arr2 の末尾に追加する
  • arr2 から重複する要素を削除する
  • arr2 のサイズを返す

Pythonによる実装例

理解を深めるために、以下の実装を見てみましょう。

def solve(c, m, n):
    if (n == 1):
        return 0
    else:
        temp_arr=[0 for i in range(2 * n+2)]
        temp_arr[0] = 0
        temp_arr[1] = c
        arr2 = []
        for i in range(2, 2 * n+2):
            temp_arr[i] = (temp_arr[i - 1] + temp_arr[i - 2]) % m
        for i in range(2, 2 * n-2, 2):
            temp = (temp_arr[i] * temp_arr[i + 2] + temp_arr[i + 1] * temp_arr[i + 3]) % m
            arr2.append(temp)
            temp = (temp_arr[i] * temp_arr[i+4] + temp_arr[i+1] * temp_arr[i+5]) % m
            arr2.append(temp)
        temp = (temp_arr[2 * n-2] * temp_arr[2 * n] + temp_arr[2 * n- 1] * temp_arr[2 * n+1]) % m
        arr2.append(temp)
        arr2 = set(arr2)
        return len(arr2)

print(solve(5, 6, 4))

入力

5, 6, 4

出力

3

コードのポイント

この実装では、まず漸化式に従って必要な項まで数列を一括生成します。その後、数列中の連続する値のペアからベクトルを構成し、ベクトル同士の内積を mod m で計算してリスト arr2 に格納していきます。最後に set() に変換することで重複が自動的に除去され、len() によって異なる剰余値の個数を取得できます。数列の生成も内積の計算もともに線形時間（O(n) 程度）で処理できるため、n がある程度大きくなっても効率的に動作するのが特徴です。