関数 \(f(N)\) を「黒板に書かれた \(N\) を全て \(1\) にするのにかかる費用」と定義します。すると \(f(N)\) は

\[ f(N) = \begin{cases} 0 & (N = 1) \\ f \left( \left\lfloor \frac{N}{2} \right\rfloor \right) + f\left( \left\lceil \frac{N}{2} \right\rceil \right) + N & (N \neq 1) \end{cases} \]

という漸化式で表せて、求めたい答えは \(f(N)\) です。よって、\(N\) が十分小さい場合はこの問題は下の実装例のように再帰関数を利用すれば解くことが出来ます。しかし、この実装では \(N\) が大きい時に実行時間が極めて遅くなり TLE してしまいます。

• 素朴な実装例 (ただし TLE してしまう)

#include <iostream>
long long f(long long N) {
  return N == 1 ? 0 : f(N / 2) + f((N + 1) / 2) + N;
}
int main() {
  long long N;
  cin >> N;
  cout << f(N) << endl;
}

そこで、関数の特徴を利用して何らかの高速化が出来ないか考えてみましょう。
試しに関数の呼び出される経路を図で書いてみます。つまり、例えば \(f(43)\) の内部では \(f(22)\) と \(f(21)\) が呼び出されるので \(43\) から \(22\) と \(21\) に矢印を引いてみます。
すると下の図のようになります。(\(N = 43\) の場合)

図を見ると分かる通り、訪問する状態は非常に少ないということがわかります。具体的には、次の 2 種類の状態のみを訪問していることが分かります。

• \(43/2^n\) を小数点以下切り捨てた値 \((21, 10, 5, 2, 1)\)
• \(43/2^n\) を小数点以下切り上げた値 \((22, 11, 6, 3, 2, 1)\)

この事実は一般化することが出来て、上の実装例において関数 f(N) を呼び出したときに呼び出される状態は

• 正整数 \(n\) を用いて \(\left\lfloor \dfrac{N}{2^n} \right\rfloor\) と表せる整数　と　正整数 \(n\) を用いて \(\left\lceil \dfrac{N}{2^n} \right\rceil\) と表せる整数

の 2 種類に限られます。(証明は帰納法を利用すると可能ですが煩雑なので省略します。)

よって、再帰関数を高速化するテクニックである メモ化再帰 を利用することで上の実装例を高速化することが出来ます。

メモ化再帰について説明します。関数 \(f(N)\) について、 \(N\) をキー、 \(f(N)\) の返り値を値とした辞書(連想配列)を作って、 \(f(N)\) を初めて計算するたびに返り値を保存していくことを考えましょう。そして、2 回目以降の呼び出しでは、関数にある処理を行わずに辞書に載っている値を返すようにします。このような処理を挟むことで高速化する手法をメモ化再帰と呼びます。( ABC247C の解説 の説明を援用しました。)

メモ化再帰を行うことで引数が同じ値の関数を無駄に何回も呼び出すことが無くなるため、計算量が大幅に改善されます。実際、メモ化再帰を行わない場合の計算量は \(\mathrm{O}(N)\) なのに対して、メモ化再帰を行うと \(\mathrm{O}(\log N)\) や \(\mathrm{O}(\log^2 N)\) 程度に高速化することが出来ます。

よって、メモ化再帰を利用してこの問題を高速に解くことが出来ます。実装例は次の通りです。

• 実装例(C++)
  • C++ では std::map のような辞書型を利用することで再帰関数をメモ化することが出来ます。

#include <iostream>
#include <map>
using namespace std;
map<long long, long long> m;
long long f(long long N) {
  if (N == 1) return 0;
  if (m.count(N)) return m[N];
  return m[N] = f(N / 2) + f((N + 1) / 2) + N;
}
int main() {
  long long N;
  cin >> N;
  cout << f(N) << endl;
}

• 実装例(Python)
  • Python では、実装例のように @cache というデコレータを付けることで自動的に再帰関数がメモ化されます。(Python 3.9 以降で登場した機能で、それ以前に使われていた @lru_cache よりも使い勝手が良いです。)

from functools import cache

@cache
def f(N):
    return 0 if N == 1 else f(N // 2) + f((N + 1) // 2) + N

print(f(int(input())))