Kazun の競プロ記録

競技プログラミングに関する様々な話題を執筆します.

AtCoder Beginner Contest 249 D問題 Index Trio

AtCoder ABC ABC249 D問題 400 pts

問題

atcoder.jp

提出解答

解法1 atcoder.jp

解法2 atcoder.jp

問題の概要

長さ N の整数列 A=\left(A_i \right)_{i=1}^N が与えられる. 次を満たす整数のトリプル *1 (i,j,k) を求めよ.

  • 1 \leq i,j,k \leq N
  • \dfrac{A_i}{A_j}=A_k

制約

  • 1 \leq N \leq 2 \times 10^5
  • 1 \leq A_i \leq 2 \times 10^5

解法

\chi_A(x) で A にある x の個数を表すとする. また, \displaystyle \dfrac{A_i}{A_j}=A_k \iff A_i=A_j A_k なので, 求めるべきは

\displaystyle \sum_{\substack{1 \leq a,b,c \leq \max A\\ a=bc}} \chi_A(a) \chi_A(b) \chi_A(c)

である.

従って, 1 \leq a,b,c \leq \max(A), a=bc を満たすトリプル (a,b,c) を求める問題に帰着された. これには2つの解決方法がある. なお, M:=\max A とする.

解法1 ( a を固定する)

a=\alpha と固定する. このとき, b は \alpha の約数でなくてはならない. また, a,b が決定されると, c=\dfrac{a}{b} と一意的に決定する. \alpha の約数列挙は計算量 O(\sqrt{\alpha}) でできることから, \alpha=1,2, \dots, M とすることにより, 時間計算量 O(M \sqrt{M}) で求めることができた.

解法2 ( b (または c) を固定する)

b=\beta と固定する. このとき, c=1,2, \dots, \left \lfloor \dfrac{M}{\beta} \right \rfloor それぞれに対して, 1 \leq a=bc \leq M であり, 1 \leq a=bc \leq M となるような c はこれで全てである. よって, 各 \beta に対して,

(\beta c, \beta, c) \quad \left(1 \leq c \leq \left \lfloor \dfrac{M}{\beta} \right \rfloor \right)

というトリプルを生成すれば良い.

計算量について, b=\beta のときに取りうる c の個数は高々 \dfrac{M}{\beta} である. よって,

\displaystyle \begin{align} \sum_{\beta=1}^M \dfrac{M}{\beta} &\leq M+\int_1^M \dfrac{M}{x} dx \\ &=M+M\left[\log |x| \right]_{x=1}^M \\ &=M+M \log M \\ &=O(M \log M) \end{align}

である.

以上から, 解法 1 では時間計算量 O(N+M \sqrt{M}), 解法 2 では時間計算量 O(N+M \log M) で求めることができた *2.

*1:三つ組のこと

*2:解法 2 では無駄なくトリプルを生成できるため, 高速になっている