【AtCoder】ABC 414 E - Count A%B=C | 緑コーダーが解くAtCoder

atcoder.jp

配点: 475 点 / 実行時間制限: 2 sec / メモリ制限: 1024 MB / Difficulty: 1296 / NoviSteps: ???

問題概要
制約
考察
実装例
コメント

問題概要

整数の組 $(a, b, c)$ であって、以下の条件を満たすものの個数を $998244353$ で割った余りを求めよ。

$1 \leq a,b,c \leq N$
$a,b,c$ は相異なる。
$a$ を $b$ で割った余りは $c$ に等しい。

制約

$3 \leq N \leq 10^{12}$
$N$ は整数

考察

問題文の言い換え

まず、3つ目の条件について考えてみる。これは以下のように言い換えられる。

$a = bx + c \quad (x は整数で, x \geq 1 を満たす)$

もし $x = 0$ だと $a = c$ になってしまい、これは2つ目の条件に反する。

また、商と余りの関係から、1つ目の条件と合わせて $1 \leq c \lt b$ が言える。

$b, c, x$ は全て正の整数なので、 $c \lt b \lt a$ という大小関係になる。したがって、2つ目の条件は常に満たされることになる。

$c$ は $a, b$ から自動的に求まることも踏まえると、この問題を以下のように言い換えることができた。変数が1個減った！

整数の組であって、以下の条件を満たすものの個数をで割った余りを求めよ。
- $1 \leq b \lt a \leq N$
- $a$ は $b$ の倍数ではない。

個数の数え上げ

では、上述の条件を満たす整数組 $(a, b)$ を数え上げていこう。

とりあえず $b = B$ と固定すると、 $B \lt a \leq N$ を満たす整数 $a$ で $B$ の倍数となるものは、

$\begin{align*} B \lt Bx \leq N & \iff 1 \lt x \leq \dfrac{N}{B} \\ \rightarrow x & = 2, 3, \dots, \left\lfloor \dfrac{N}{B} \right\rfloor \end{align*}$

より、 $\left\lfloor \dfrac{N}{B} \right\rfloor - 1$ 個存在する。したがって、 $B$ の倍数でない $a$ は $(N - b) - \left( \left\lfloor \dfrac{N}{B} \right\rfloor - 1 \right)$ 個ある。

これを全ての $B$ について和を取ればよいので、求める答えは

$\begin{align*} \displaystyle \sum_{b=1}^{N} \left( N - b + 1 - \left\lfloor \dfrac{N}{b} \right\rfloor \right) \end{align*}$

と書ける。

答えを高速に求める

さて、上の式を愚直に求めていては $O(N)$ の計算量となってしまい、本問の制約下では TLE になる。何とかして高速化しなければならない。

とりあえず、答えの式を以下のように分解してみよう。

$\begin{align*} \displaystyle \sum_{b=1}^{N} \left( N - b + 1 \right) - \sum_{b=1}^{N} \left\lfloor \dfrac{N}{b} \right\rfloor \end{align*}$

第1項に関しては、等差数列の和の公式から $\frac{N(N+1)}{2}$ と $O(1)$ で計算できる。普通にシグマを外してもよい。

問題は第2項だ。

実は、このような「分数の floor sum を求めろ」という問題が ABC230-E で出題されている。

atcoder.jp

いくつか解法はあるが、結局この部分は $O(\sqrt{N})$ で以下のように計算できる(商列挙というやつ)。 $b = \sqrt{N}$ で2つに分けて、対称性を利用する。

$\begin{align*} \sum_{b=1}^{N} \left\lfloor \dfrac{N}{b} \right\rfloor = 2 \sum_{b = 1}^{\lfloor \sqrt{N} \rfloor} \left\lfloor \dfrac{N}{b} \right\rfloor - \lfloor \sqrt{N} \rfloor^2 \end{align*}$

以上より、求める答えは $O(\sqrt{N})$ で計算することができた。

実装例

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

#if __has_include(<atcoder/all>)
#include <atcoder/all>
using namespace atcoder;
#endif

#define repe(i, start, end) for (auto i = (start); (i) <= (end); (i)++)

using ll = long long;
using mint = modint998244353;

// ======================================== //

int main()
{
    ll N;
    cin >> N;

    mint sum = 0;
    ll sqrt_N = floor(sqrt(N));
    repe(b, 1, sqrt_N) sum += N / b;

    mint m_sqrt_N = sqrt_N;
    mint ans = (mint)N * (N + 1) / 2 - 2 * sum + m_sqrt_N * m_sqrt_N;

    cout << ans.val() << endl;
    return 0;
}