Division or Subtractionを解く(AtCoder Beginner Contest 161 F)

導入

AtCoder Beginner Contest 161 F - Division or Subtractionを解きます。

解法については公式解説が存在します。 この記事では私が解法に至るまでの思考を書きます。 思考の過程を解説するので、やや遠回りもします。 どうやって解法を思いつくかの参考になれば幸いです。

前提

問題文はAtCoderで確認してください。

問題における、「最終的にNが1になるようなK」を本文では便宜上「いいK」と呼びます。

本文

Nが大きい

問題文を読んだ時にまず注目したのは、 Nの大きさです。 2 <= N <= 10 ^ 12 から 2 <= K <= N 全てについてKがいいKかどうかO(1)で調べられたとしても実行時間制限に間に合わないということです。

# このような解法は check がどれだけ早くても間に合わない
ans = 0
for k in range(2, n+1):
  if check(n, k):
    ans += 1
print(ans)

なので、TLEにならないためにはたとえば下記のようなことができないかを考える必要があるということを頭に入れておきます。

  • いいKになりうる候補をO(√n)個以下程度に絞り込む(O(log n)などだとより嬉しい)。
    • たとえば√n個に絞り込めた場合checkO(log n)程度で実装できれば候補に対してcheckしていく方法が間に合う。
    • あるいは、log n個に絞り込めた場合checkO(√n)程度かかっても間に合う。
  • Kに対していいKかを調べる方法ではなく、NからいきなりいいKの個数をO(√n)以下で求める。
  • Nに対しいいKM個以上あるか」をO(√n)以下で求める。
    • これができるなら二分探索で答えを求めることができる

操作について

操作の性質について考えていきます。

一旦、手で操作を実施してみます。人によるかもしれませんが、具体的な値を触った方が理解が深まるということはよくあります。 入力例1にあるN=6を使ってみます。

  • K=2のとき(出力例の説明にあるとおり、これは「いいK」のはず)
    1. 6は2で割り切れるので割ります。6→6/2=3
    2. 3は2で割り切れないので引きます。3→3-2=1
    3. 1はK=2未満なので終了です。いいKです。
  • K=3のとき(出力例の3つの「いいK」いずれでもなく、「いいK」ではないはず)
    1. 6は3で割り切れるので割ります。6→6/3=2
    2. 2はK=3未満なので終了です。いいKではありません。

操作をプログラムで行なった場合どれくらいの時間がかかるかを見積もってみます。 条件によって2種類の操作を使い分けて操作を繰り返した場合の計算時間の見積もりは少し難しいので、一旦片方の操作しかない場合を考えてみます。

  • 割り算だけを考える場合
    • 簡単のためNK未満になるまで必ず割り切れる場合を考えます。一回の操作でN1/Kの大きさになります。これは指数関数的な減少であり、NK未満になるまでlog_k n回程度の操作で済みそうです。
  • 引き算だけを考える場合
    • 一回の操作でNKだけ小さくなるので、繰り返しにより実現するとn/k回の操作が必要です。しかし、もし引き算のみを繰り返す場合は割り算の余りと同じですのでn%kとすることでO(1)で最終的なNを求めることができます。

単一の操作の繰り返しはどうやらある程度高速に実行できそうです。 あるいは、もし実はほとんどどちらかの操作の連続でごくまれに操作が切り替わるのであれば高速に操作を終わらせることができるかもしれません。

操作が連続しそうか、どんなときに切り替わるかについて考えてみます。

  • 割り算を実施した後
    • 割り算をした後再び割り切れるか、あるいは何連続で割り切れるかを高速で判断する方法は残念ながら思いつきませんでした。
    • 余談ですが、Kが固定であれば[1, K, K^2, K^3,...]というリストをあらかじめ作っておくことで二分探索で何連続割り切れるか高速に判断できそうです。今回はいろんなKについて調べたいので不適です。
  • 引き算を実施した後
    • 引き算を実施したということは引き算実施前のNKで割り切れません。この場合、N-KKで割り切れないので、引き算の後の操作は必ず引き算ということがいえます。

上記考察により引き算のあとに割り算をすることはないことがわかりました。 このことから、KがいいKかどうか調べるのは次のような手順で実現できます。

def check(original_n, k):
  n = original_n
  while n % k == 0:
    n = n // k
  n = n % k
  return n == 1

この操作は、繰り返し文の箇所が高々log_k n回程度しか行われないので、もしいいKになりうる候補をO(√n)個以下に絞り込めればそれらに対してチェックをしても間に合いそうです。

候補の絞り込み

候補の絞り込みができるか考えてみます。 この手順で解けると決まったわけではありませんが、他に確実な解法も思いつかないので一旦考えてみるという感じです。

これまでの考察から、操作は0回以上の割り算→0回以上の引き算の順しかないので、Kを2種類に分類して考えてみます。

  • 1回以上の割り算を行うもの = Nの約数
  • 1度も割り算を行わないもの = Nの約数ではない

やや方針が突飛に見えるかもしれませんが、下記の点からこの分類は考察する価値がありそうです。

  • Nの約数の数は高々2√n個なので、前者は絞り込みの要件に合う
  • 後者は操作が1つしかないので考察を進めやすそう

まず、Nの約数であるKについては個数が2√N個以下であり、その列挙もO(√n)であることがいえます。

# 約数を列挙する方法の例
def divisors(n):
  result = []
  for i in range(1, n+1):
    if i*i > n:
      break
    if n%i == 0:
      result.append(i)
      d = n // i
      if d != i:
        result.append(d)
  return result

これら全てに対していいKかどうか調べることは実行時間制限内に収まりそうです。

Nの約数ではないKについて考察します。 引き算のみを行うので、N % K == 1となるようなKの数を調べるということになります。 この変形ができてしまえば簡単で、N-1の約数がいいKということになります。 N-1の約数もNの約数と同様O(√n)で列挙できるので、どうやら間に合いそうです。

仕上げ

Nの約数であるようないいKの数とNの約数でないようないいKの数の合算が答えです。

次のことに気をつけます。

  • Kは2以上なのでNN-1の約数のうち「1」はカウントしない
  • 同じKを重複してカウントしてはいけないが、NN-1の最大公約数は1なので考慮は不要

以上より、下記のような手順で解くことができます。

  1. Nの約数であるようないいKをカウントする
    1. 1以外のNの約数を全て列挙する
    2. 各約数について、いいKかチェックする
  2. Nの約数でないようないいKをカウントする
    1. 1以外のN-1の約数の数がNの約数でないようないいKの数
  3. Nの約数であるようないいKの数とNの約数でないようないいKの数の和が答え

私の提出