ARC174D - Digit vs Square Root（yを埋め込む解法）

TL;DR

$y$ を固定したとき、対応する $x$ の個数が $\mathcal{O} ( \log y )$ で求められる。
$1 \leq y \leq 10^9$ の範囲を調べればよく、手元なら何とか前計算できる。
前計算すると、考慮すべき（i.e. 対応する $x$ が 1個以上ある） $y$ はかなり少ないことが分かるので、それを埋め込めば通る。

TL;DR
問題概要
考察
- 実験
  - 実験コード
  - 出力
- y を固定したときの x の個数
  - 条件を区間に変形する
前計算・埋め込み
- 前計算
  - 前計算コード
  - 出力結果
提出
- 提出用コード
- 提出結果
あとがき

問題概要

以下の条件をすべて満たすの個数を求めてください。
- $x$ は $N$ 以下の正整数である。
- $\lfloor \sqrt{x} \rfloor$ （ $\sqrt{x}$ の整数部分）を $y$ とするとき、 $y$ は $x$ の接頭辞になる。
$T$ ( $1 \leq T \leq 10^5$ ) ケース与えられるので、それぞれに答えてください。

考察

実験

例を見ると、 $x=1$ や $x=100$ は条件を満たすことが分かる。 $x=10^{2n}$ の形で表される $x$ は条件を満たすっぽい。それ以外で条件を満たす $x$ にはどんなものがあるか、愚直に実験してみる。公式解説とほぼ同様の実験コードだが、 $x$ に対応する $y$ を出力するようにしてある。

実験コード

#include<bits/stdc++.h>

using namespace std;

long long sq(long long x){
  long long l=0,r=2e9;
  while(l<=r){
    long long te=(l+r)/2;
    if(te*te <= x){l=te+1;}
    else{r=te-1;}
  }
  return r;
}

int main(){
  for(long long x=1;x<=11111;x++){
    long long y=sq(x);
    string s=to_string(x);
    string t=to_string(y);
    bool ok=true;
    for(int i=0;i<t.size();i++){
      if(s[i]!=t[i]){ok=false;}
    }
    if(ok){cout << x << " " << y << "\n";}
  }
  return 0;
}

出力

出力結果を眺めていると $x$ にも法則性がありそう（公式解説はこの法則性を使って解いている）だが、 $y$ としてありうる値はかなり限られていそうなことが分かる。 $y$ が少ないならば、 $y$ から $x$ の個数を逆算して速く処理できないだろうか？

y を固定したときの x の個数

ある $y$ を固定したとき、題意の条件を満たす $x$ ( $\leq N$ )の個数を $f(y, N)$ とする。

いったん実験結果を忘れて、 $f(y, N)$ を求める方法を考えてみる。

条件を区間に変形する

問題概要で $x$ が満たすべき条件は以下のようだった（ $x \leq N$ という条件もとりあえず無視）。

$\lfloor \sqrt{x} \rfloor$ （ $\sqrt{x}$ の整数部分）を $y$ とするとき、 $y$ は $x$ の接頭辞になる。

$y$ を固定することで、これが2つの条件で言い換えられることになる。

$\lfloor \sqrt{x} \rfloor = y$
$y$ は $x$ の接頭辞である。

さらに、それぞれの条件を $x$ の不等式でこのように書き換えることができる。

$y^2 \leq x \lt ( y + 1 ) ^ 2$
$10^k y \leq x \lt 10^k ( y + 1 )$ となる非負整数 $k$ が存在する。

2つ目の条件の書き換えはやや非自明だが、次のように考えることができる。

例えば $y=12345$ のとき¹、2つ目の条件を満たす $x$ は $x = 12345 ?? \dots ?$ （ $?$ は任意の数字）の形になる。 $?? \dots ?$ の部分を $c$ で表して、その桁数を $k$ とする²と、 $x = 12345 \times 10 ^ k + c$ となる。ここで $c$ が $12345$ の部分に影響を及ぼさないようにするためには $0 \leq c \lt 10 ^ k$ でなければならず、この不等式を変形すると結局上の式になる。

この形で表せば、 $f(y, N)$ を求める問題は、不等式で表される2つの区間の共通部分を求める問題に帰着できる。新たな変数 $k$ が導入されてしまったが、2つの区間が共通部分を持つのは上の不等式で表される区間の右端と下の不等式で表される区間の左端が $10^k y \lt ( y + 1 ) ^ 2$ を満たすときに限られる。つまり、 $k$ としては $\mathcal{O} ( \log y )$ 個ぐらいの値を調べればよいだけなので、計算量的には問題ないだろう。また、区間で表されているので、無視していた $x \leq N$ という条件を再度取り込むのにも相性がよさそうだ。

前計算・埋め込み

$1 \leq y \leq \sqrt{N}$ の範囲の $y$ について $f(y, N)$ を足し合わせれば答えが求まるが、計算量が1ケースあたり $\mathcal{O} ( \sqrt{N} )$ となり、そのまま提出すると間に合わない。だが、実験結果から、実際にありうる $y$ はかなり限られていそうだったので、手元でありうる $y$ を前計算しておけば何とか通せないだろうか？

前計算

答えを求めるには $f(y, N) \gt 0$ となりうる $y$ だけを考えればよい。すなわち、 $x \leq N$ の条件を考慮せず $N = \infty$ としても $f(y, \infty) = 0$ となるのであれば、提出用コードではその $y$ を調べる必要はない。

前計算コードでは $1 \leq y \leq \sqrt{10 ^ {18}}$ の範囲を全探索して、 $f(y, \infty) \gt 0$ となる $y$ を出力している。提出用コードではこの $y$ だけを調べればよい。

前計算コード

#include<bits/stdc++.h>
using namespace std;
using ll = long long;
const ll ymax = 1'000'000'000LL;

inline ll div_ceil(ll num, ll den){
    // 切り上げ除算
    return (num + den - 1) / den;
}

ll f(ll y){
    // y を決め打ちしたときの対応する x の個数
    // x <= N の条件は考慮しない

    // あらかじめ 10^k の上限（区間が共通部分を持つ条件）を求めておく
    // （オーバーフロー対策）
    ll b_lim = div_ceil((y + 1) * (y + 1), y);

    ll ret = 0;
    for(ll b = 1; b < b_lim; b *= 10){
        // b = 10^k. 各 k に対して x の個数を求める
        ll l = max(y * y, b * y); // 2つの区間の下限をまとめる
        ll r = min((y + 1) * (y + 1), b * (y + 1)); // 2つの区間の上限をまとめる
        ret += max(0LL, r - l);
    }
    return ret;
}


int main(){
    // 提出用コードに埋め込むべき y を出力する
    for(ll y = 1; y <= ymax; y++){
        if(f(y) > 0) cout << y << ", "; // f(y, ∞) > 0 であれば出力
    }
    cout << endl;
    return 0;
}

出力結果

30秒ぐらいで出てきました（i9-10850K)。

1, 8, 9, 10, 98, 99, 100, 998, 999, 1000, 9998, 9999, 10000, 99998, 99999, 100000, 999998, 999999, 1000000, 9999998, 9999999, 10000000, 99999998, 99999999, 100000000, 999999998, 999999999, 1000000000,

$y$ としてありうる値は少ないとは思ってたけど、びっくりするぐらい少なかった。

提出

先ほど求めておいた $y$ を埋め込んで提出する。前計算用コードを少しいじれば提出用コードになる。

提出用コード

#include<bits/stdc++.h>
using namespace std;
using ll = long long;
const ll ymax = 1'000'000'000LL;

inline ll div_ceil(ll num, ll den){
    // 切り上げ除算
    return (num + den - 1) / den;
}

ll f(ll y, ll N){
    // y を決め打ちしたときの対応する x (<= N) の個数

    // あらかじめ 10^k の上限（区間が共通部分を持つ条件）を求めておく
    // （オーバーフロー対策）
    ll b_lim = div_ceil((y + 1) * (y + 1), y);

    ll ret = 0;
    for(ll b = 1; b < b_lim; b *= 10){
        // b = 10^k. 各 k に対して x の個数を求める
        ll l = max(y * y, b * y); // 2つの区間の下限をまとめる
        ll r = min({(y + 1) * (y + 1), b * (y + 1), N + 1}); // 2つの区間の上限 && x <= N をまとめる
        ret += max(0LL, r - l);
    }
    return ret;
}

ll solve(ll N){
    // 各ケースの問題を解く

    const vector<ll> ys({1, 8, 9, 10, 98, 99, 100, 998, 999, 1000, 9998, 9999, 10000, 99998, 99999, 100000, 999998, 999999, 1000000, 9999998, 9999999, 10000000, 99999998, 99999999, 100000000, 999999998, 999999999, 1000000000}); // 埋め込み
    ll ret = 0;
    for(ll y : ys) ret += f(y, N);
    return ret;
}

int main(){
    int T;
    cin >> T;
    for(int t = 0; t < T; t++){
        ll N;
        cin >> N;
        cout << solve(N) << endl;
    }
    return 0;
}

提出結果

AC. はいキン肉マンおめでとう

https://atcoder.jp/contests/arc174/submissions/51430684

あとがき

コンテスト中は $x$ の満たす性質を実験→未証明ACしていた人が多かったようだ。この記事でも手計算での証明は行っていないが、ありうる $y$ を全探索して埋め込んでいるので一応の正当性は確保できているものと思いたい……

実際には $y=12345$ のときには対応する $x$ は存在しないが、説明の便宜上この値とした。↩
$?? \dots ?$ の部分だけを抜き出したときに leading zero(s) がある場合も桁数に数える。↩