Pythonで ax + by = n が解を持つ係数ペア (a, b) の個数を求めるプログラム

問題の概要

ある整数 n が与えられたとき、方程式 a*x + b*y = n が少なくとも1つの解を持つようなペア (a, b)(ただし a < b)の個数を求めることを考えます。

例えば、入力が n = 4 の場合、条件を満たすペアは (1, 2) と (1, 3) の2つであるため、出力は 2 になります。

解法のアプローチ

この問題を解くために、以下の手順に従います。

• 引数として n を受け取る関数 divisors_gen() を定義します。
• divs := サイズ n+1 のリストのリストとし、各内部リストには 1 を格納します。
• divs[0] := 要素 0 のみを持つリストとします。
• i を 2 から n まで繰り返します。
  • j を 1 から ⌊n / i⌋ + 1 まで繰り返します。
    • インデックス [i * j] のリストの末尾に i を追加します。
• すべての内部リストを反転した状態で divs を返します。

メイン処理では、以下の手順を実行します。

• result := 0 と初期化します。
• d_cache := divisors_gen(n+1) として約数リストを事前にキャッシュします。
• a を 1 から n - 1 まで繰り返します。
  • i := 1 とし、s := 新しい空の集合を用意します。
  • a*i < n の間、以下を繰り返します。
    • b := n - a*i とします。
    • d_cache[b] 内の各 d について:
      • d > a の場合、d が s に存在しなければ result を 1 増やします。
      • それ以外の場合はループを抜けます。
      • d を集合 s に追加します。
    • i を 1 増やします。
• 最後に result を返します。

アルゴリズムのポイント

このアルゴリズムでは、篩(ふるい)法に似た手法を使って各数値の約数リストを効率的に生成しています。その後、a を固定した際に b = n - a*i となる各候補に対して、b の約数の中から a より大きいものだけを数え上げます。集合 s を使って既にカウントした約数を記録することで、同じペアの二重計上を防ぎ、重複なく正確にペア (a, b) を数えることができます。

実装例

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

def divisors_gen(n):
   divs = [[1] for x in range(0, n + 1)]
   divs[0] = [0]
   for i in range(2, n + 1):
      for j in range(1, n // i + 1):
         divs[i * j].append(i)
   return [i[::-1] for i in divs]

def solve(n):
   result = 0
   d_cache = divisors_gen(n+1)

   for a in range(1, n):
      i = 1
      s = set([])
      while a*i < n:
         b = n - a*i
         for d in d_cache[b]:
            if d > a:
               if d not in s:
                  result += 1
            else:
               break
            s.add(d)
         i += 1
   return result

n = 4
print(solve(n))

入力

4

出力

2