ABC292 C問題（Four Variables）を解く

AtCoder が提供しているABC（AtCoder Beginner Contest）292 のC問題をC++とPythonで解いてみました。ABC292は、2023年3月4日21:00に実施されました。

AtCoder の紹介はこちらに、プログラミングの方針はこちらに記事があります。

C問題 Four Variables（Difficulty : 444）

問題はリンク先をご覧ください。

ABC292 C問題 Four Variables

計算時間を短くする工夫が必要な問題です。AtCoder Problems による Difficulty は 444 でした。

計算量の考察

問題の制約では $2 \leqq N \leqq 2 \times 10^5$ となっています。環境と演算に依存しますが、1秒間にできる演算は、10⁸ または 10⁹ 回程度となります。

以下の4重ループは、素朴な全検索ですが、明らかに間に合いません。入力例1（N = 4）は正解しますが、入力例2（N = 292）では間に合いません。

	for (int a = 1; a < n; ++a) {
		for (int b = 1; b < n; ++b) {
			for (int c = 1; c < n; ++c) {
				for (int d = 1; d < n; ++d) {
					if (a * b + c * d == n) {
						++result;
					}
				}
			}
		}
	}

A、B、C が決まれば、D が決まります。D = (N – A × B) / C が割り切れれば、条件を満たします。結果的に計算量を3重ループにすることができました。各ループ変数の上限も可能な範囲で少なくしています。以下が、3重ループのプログラムの抜粋です。

	for (int a = 1; a < n; ++a) {
		for (int b = 1; a * b < n; ++b) {
			for (int c = 1; c <= n - a * b; ++c) {
				if ((n - a * b) % c == 0) {
					++result;
				}
			}
		}
	}

ただし、このコードも入力例2（N = 292）は正解しますが、入力例3（N = 19876）では間に合いません。

A × B と C × D は同じ計算をすることになります。A × B の結果（その積に等しいAとBの組み合わせの数）をmapに格納しておけば、C × D に相当する計算は map で値を参照することで処理することができます。

map に格納するコードは以下となります。

	map<int, int> m;
	for (int a = 1; a < n; ++a) {
		for (int b = 1; a * b < n; ++b) {
			++m[a * b];
		}
	}

この工夫で間に合うように計算量を下げることができました。

解答案

C++ プログラム例（ABC292C）

上記で考察した工夫でプログラムします。ABの積について通り数を map に格納して、CD = N – AB の値が map に格納されているか調べています。格納されている場合は、AB の通り数と CD の通り数を掛けた値を解答に加えます。

解答 result が32ビット整数に収まるのか不明であったため、64ビット整数にしました。ただし結果的には、N＝2 × 10⁵ の場合も、解答は、32ビット整数の範囲に収まっていました。

以下が、C＋＋プログラムとなります。

#include <bits/stdc++.h>
using namespace std;

typedef unsigned long long int ull;

int main()
{
	int n;
	cin >> n;

	map<int, int> m;
	for (int a = 1; a < n; ++a) {
		for (int b = 1; a * b < n; ++b) {
			++m[a * b];
		}
	}

	ull result = 0;
	for (auto it = m.begin(); it != m.end(); ++it) {
		int cd = n - (*it).first;
		if (m[cd] != 0) {
			result += (*it).second * m[cd];
		}
	}

	cout << result << endl;

	return 0;
}

AC（Accepted＝正しいプログラム）と判定されました。

Python プログラム例（ABC292C）

Python は、defaultdict を使いました。C++ のプログラムを移植しただけなので、Python らしいプログラムではないかもしれません。

"""AtCoder Beginner Contest 292 C"""
from collections import defaultdict

n = int(input())
m = defaultdict(int)
for a in range(1, n):
    b = 1
    while a * b < n:
        m[a * b] += 1
        b += 1

result = 0
for k, v in m.items():
    cd = n - k
    if cd in m:
        result += v * m[cd]

print(result)

こちらも「AC」と判定されました。

最後に

3重ループの解答は思いつきましたが、提出するまでもなく計算時間が間に合っていません。そのため、D問題とE問題に先に取り組みました。結果的に、E問題は解けませんでした。残り時間10分でC問題に戻りmapを使う実装を思いつきました。AC判定を得たのが、99分26秒でぎりぎりでした。

ここまで粘るのは珍しく、AC判定をうれしく感じました。

引き続き ABC の問題を紹介していきます。