ARC060 D - 桁和 / Digit Sum

問題

arc060.contest.atcoder.jp

考察

まず、nをb進数で表した際、各桁の合計がsになるかどうかは簡単に求まります。
一番時間がかかる2進数ですら、64回ループするだけで[tex: 1018]程度までカバーできるので。
(問題分の定義では再帰で書かれていますが、whlieループに書き換えておくとより速くていい感じです)

さて、nをb進表記するとはどういうことなのでしょうか。
1, b, b2, b3, ... という要素で割った値に分割するということと捉えます(そのままですが)。
つまり、nがbで割り切れる時、各桁の和は1以上になります。
nが大きいですが、この方法で 2 \leq b \leq {\sqrt{n}}までチェックできます。
条件を満たす最小のものが見つかれば、それが答えです。

これより大きいbはどのようにチェックすれば良いのでしょうか?
上記のケースに当てはまらない場合というのは、存在すれば1, bの桁のみで構成されます。
なぜならば、b2はnより大きくなってしまうためです。

1の桁の数をp、bの桁の数をqとおいて考えてみると、以下の条件が成り立ちます。

  1.  p + b * q = n
  2.  p + q = s
  3.  1 \leq q < {\sqrt{n}} (繰り上がってしまうため)

1,2の式を使うと、 b = {\frac{n-s}{p}} + 1となり、pを全探索すればbが求まります。
なお、 n - s < 0 となると b = 0になるため不適切です。

最後に、例外ケースを考慮し、n=sの時はb=n+1にするだけで良いです。

コード

ll dig(ll b, ll n) {
    ll ans = 0;
    while(n > 0) {
        if(n < b) {
            ans += n;
            break;
        }
        ans += (n % b);
        n /= b;
    }
    return ans;
}

int main() {
    ll n, s;
    cin >> n >> s;

    if(n == s) {
        cout << n+1 << endl;
        return 0;
    }

    for(ll i = 2; i*i <= n; i++) {
        if(dig(i, n) == s) {
            cout << i << endl;
            return 0;
        }
    }

    ll ans = 1ll<<60;
    for(ll i = 1; i*i < n; i++) {
        if( n - s >= 0 && (n-s) % i == 0) {
            ll b = ((n-s) / i) + 1;
            if(dig(b, n) == s)
                chmin(ans, b);
        }
    }

    if(ans != 1ll<<60)
        cout << ans << endl;
    else if(s == 1)
        cout << n << endl;
    else
        cout << -1 << endl;

    return 0;
}