mayoko’s diary

プロコンとかいろいろ。

Codeforces Round #332 (Div. 2) D. Spongebob and Squares

こういう問題でギリギリ攻められるのあんまり好きじゃないんですが…
まぁこの問題の場合だとまだ仕方ないかなって感じしますけど…

解法

m < n とすると, m*n のテーブルの中にある正方形の数は
m*n + (m-1)*(n-1) + ... + 1 * (n-m+1)
となります。そこで, m の値を決め打ちして, n の値が存在するかどうかを調べる, ということをやってみます。後でわかりますが, 上の値はおよそ m の 3 乗に比例するので, 調べなければならない数の範囲は 10^6 程度で済みます。

上の式で, 1 * (n-m+1) となっているところの, n-m+1 を a という未知数にして考えると, 正方形の数は
 1*a + 2*(a+1) + ... + m*(a+m-1) = \sum_{k=0}^m k*(a+k-1)
となるので, 適当にシグマ計算すると, O(1) で 各 a に対する正方形の数が求められます。なので, a に関して二分探索すれば良いでしょう。

計算ではオーバーフローを起こさないように注意です。工夫したのは, 二分探索の最大値は 2*1e18 / (i*i) で抑えられるので, high の値を最初からこの値にした という点ですね。C++ に多倍長欲しい…

const ll MAX = 1800000;

int main() {
    cin.tie(0);
    ios::sync_with_stdio(false);
    ll x;
    cin >> x;
    vector<pair<ll, ll> > ans;
    for (ll i = 1; i <= MAX; i++) {
        ll low = 0, high = (1ll<<61) / (i*i) + 1;
        while (high - low > 1) {
            ll med = (high+low) / 2;
            ll result = i*(i+1)/2 * (2*i+1) / 3;
            result += i*(i+1)/2 * (med-1);
            if (result <= x) low = med;
            else high = med;
        }
        ll result = i*(i+1)/2 * (2*i+1) / 3;
        result += i*(i+1)/2 * (low-1);
        if (result == x) {
            ans.emplace_back(i, low+i-1);
            ans.emplace_back(low+i-1, i);
        }
    }
    sort(ans.begin(), ans.end());
    ans.erase(unique(ans.begin(), ans.end()), ans.end());
    cout << ans.size() << endl;
    for (auto p : ans) cout << p.first << " " << p.second << endl;
    return 0;
}