AtCoder が提供しているABC(AtCoder Beginner Contest)340 のC問題をC++とPythonで解いてみました。ABC340は、2024年2月10日21:00に実施されました。
AtCoder の紹介はこちらに、プログラミングの方針はこちらに記事があります。
C問題 Divide and Divide(Difficulty : 528)
問題はリンク先をご覧ください。
メモ化再帰を使います。AtCoder Problems による Difficulty は 528 でした。
解答案
金額の総和を返す関数 f は、以下のように再帰的に定義できます。
$$f(n) = n + f \left(\left\lfloor \frac{n}{2} \right\rfloor \right) + f \left(\left\lceil \frac{n}{2} \right\rceil \right)\;\;(i \geqq 2)$$
$$f(n) = 0\;\;(i \leqq 1)$$
再帰的に同じ数が繰り返し呼ばれます。一度、呼ばれた値を記憶する「メモ化再帰」を使って計算量を下げることができます。また、n の値は半分になっていくため、$N$ の上限が $10^{17}$ と大きいですが、$10^{18}$ を $2^{60}$ で近似すると、呼び出し回数もそれほど大きくなりません。
C++ プログラム例(ABC340C)
値の幅が大きいため、メモには map コンテナを使います。漸化式はそのまま実装します。C++の割り算は切り捨てですが、$(n + 1) / 2$ として切り上げることができます。
以下が、C++プログラムとなります。
#include <bits/stdc++.h>
using namespace std;
typedef unsigned long long int ull;
map<ull, ull> m;
ull f(ull n)
{
if (n <= 1) {
return 0;
} else if (m[n] == 0) {
m[n] = n + f(n / 2) + f((n + 1) / 2);
}
return m[n];
}
int main()
{
ull n;
cin >> n;
cout << f(n) << endl;
return 0;
}
同じような考えですが、map コンテナで、黒板に書かれている数字とその回数を記憶しても解くことができます。
- map コンテナに、黒板に書かれている数字とその回数を組にして格納します。ただし、数字が2以上となる場合だけ格納します。
- map コンテナから、一番大きい数字を取り出して、処理をします。処理した値は map コンテナから削除します。
- map コンテナが空になれば、終了です。
以下がこの方針で解いたプログラムとなります。
#include <bits/stdc++.h>
using namespace std;
typedef unsigned long long int ull;
int main()
{
ull n;
cin >> n;
map<ull, ull> m;
m[n] = 1;
ull result = 0;
while (m.size() > 0) {
auto it = m.rbegin();
ull number = (*it).first;
ull count = (*it).second;
ull n1 = number / 2;
if (n1 >= 2) {
m[n1] += count;
}
ull n2 = (number + 1) / 2;
if (n2 >= 2) {
m[n2] += count;
}
m.erase(number);
result += number * count;
}
cout << result << endl;
return 0;
}
どちらも AC(Accepted=正しいプログラム)と判定されました。
Python プログラム例(ABC340C)
Python 版も基本的な考え方は同じです。メモは辞書に格納します。以下となります。
"""AtCoder Beginner Contest 340 C"""
m = {}
def f(n):
if n <= 1:
return 0
if n not in m:
m[n] = n + f(n // 2) + f((n + 1) // 2)
return m[n]
print(f(int(input())))
Python には、デコレータという文法が用意されています。ABC275 D問題でも紹介しました。Python 3.9 以降では、cache というデコレータが使えます。
以下が、デコレータを使ったプログラムです。非常にすっきりと書けています。
"""AtCoder Beginner Contest 340 C"""
import functools
@functools.cache
def f(n):
if n <= 1:
return 0
return n + f(n // 2) + f((n + 1) // 2)
print(f(int(input())))
こちらも「AC」と判定されました。
最後に
同じ趣旨の問題が、ABC275 D問題(解説記事)で出題されていました。自分で書いた記事ですが、Python のデコレータについて、すっかり忘れていました。復習ができました。
引き続き ABC の問題を紹介していきます。