【AtCoder】ABC254 のA,B,C,D,E における Python解説

ABC 254 のA,B,C,D,E問題を解くために考えたこと、ACできるPython3(PyPy3)コードを紹介します。

この記事は @u2dayo さんの記事を参考にしています。見たことのない方はそちらもご覧ください。とても勉強になります。

また、問題の難易度を表す指標を Atcoder Problems から引用しています。このサイトは勉強した問題を管理するのにとてもオススメです。

質問やご指摘はこちらまで
Twitter : Waaa1471

作者プロフィール
Atcoder ：緑色 882
20221231 現在

目次

はじめに
A.Last Two Digits
B.Practical Computing
C.K Swap
D.Together Square
E.Small d and k

はじめに

特にC問題以降になると、競技プログラミングに典型的な前提知識、少し難しい数学的考察が必要になり始めます。
しかし、公式解説ではこの部分の解説があっさりしすぎていて競技プログラミング(Atcoder)を始めたばかりの人にはわかりにくい、難しいと感じるのではないでしょうか。

またC++がわからないと、コードの書き方を勉強することが難しいです。一応、参加者全員のコードを見ること自体は可能ですが、提出コードは必ずしも教育的ではありません(ここで紹介する記事も本番で提出したものとは全く異なります)。そんなものから初学者が解説もなしになにか得ることはとても難しいと思います。実際適当に何人かのコードをみたものの、意味がわからずに終わった経験があるのではないでしょうか。

この記事がそんな方々の勉強の助けになればよいなと思っています。

A.Last Two Digits

問題

難易度：灰色 7

考察

入力を文字列型で受け取って下2桁を出力する。
もしくは、整数型で受け取って100で割った余りを出力します。ただし余りが 1桁の場合を考慮して、答えは 2桁になるようにゼロ埋めする必要があります。

コード

pypy3

文字列受け取り

N=input()
# N は3桁であるので、N[1:] でも可
print(N[-2:])

整数型受け取り

N=int(input())
print(f"{N%100:02}")

補足

• リストの任意の範囲を取りだし
  Pythonのスライスによるリストや文字列の部分選択・代入

• f-strings 書式指定
  Pythonのf文字列（フォーマット済み文字列リテラル）の使い方
  ゼロ埋め、進数変換は頻出なので出てくるたび覚えましょう。

B.Practical Computing

問題

難易度：灰色 45

考察

数式がご丁寧に与えられているので、それに従って数列を完成させていきます。具体的には $A_0$ から順に 「 直前の数列の値を用いて次の数列を完成させること 」　を繰り返して全ての数列を完成させます。

ここから 直前の数列を覚えておく必要がある ということを読み取ることができました。
実装では、直前の数列を prev リストで管理し、これを新しく作成した数列で更新することで、直前の数列を覚えている状態を維持します。

なお、各数列は空白区切りで出力する必要があります。

コード

pypy3

N=int(input())
prev=[]
for i in range(N):
    now=[]
    for j in range(i+1):
        if j==0:
            now.append(1)
        elif j==i:
            now.append(1)
        else:
            now.append(prev[j-1]+prev[j])
    print(*now)
    # 直前の数列を更新
    prev=now

全ての数列を二次元リストで管理しても OK

N=int(input())
A=[]
for i in range(N):
    now=[]
    for j in range(i+1):
        if j==0:
            now.append(1)
        elif j==i:
            now.append(1)
        else:
            now.append(A[i-1][j-1]+A[i-1][j])
    print(*now)
    A.append(now)

ちなみに

「 ある状態が遷移する次の状態を更新する 」ことを繰り返して、目的の状態を求めるアルゴリズムを 動的計画法 (DP) と言います。
遷移の中で成り立つ関係式(漸化式)に従い、状態の更新を繰り返したこの問題は動的計画法を用いていたのでした。

この問題では全ての状態を考えることができましたが、問題の難易度があがると全ての状態を考えて更新することは基本的に不可能です(計算量の観点から)。しかし、動的計画法を用いることで、考えるべき状態を取捨選択して効率的に目的の状態を求めることが可能になります。
特に D問題で出題されることがよくあるのでわからない方、興味がある方はぜひ勉強してみたください

補足

• 出力形式
  競技プログラミングにおける Python での標準入出力のまとめ
  出力形式がとても分かりやすくまとめられています。

• 動的計画法
  動的計画法ってなに？　(導入)
  典型的な DP (動的計画法) のパターンを整理 Part 1 ～ ナップサック DP 編 ～
  こちらで動的計画法の概念を勉強することができます。
  分かるようになったら、アルゴ式 1～2章で練習してみましょう
  動的計画法 (DP)

C.K Swap

問題

難易度 ： 茶色 536

考察

K>1 の場合、全体を自由に入れ替えることはできず、特定の位置でのみ入れ替えが可能です。このような限られた操作で全体を昇順に入れ替えるためには、まず入れ替え可能なグループに分類し、それらを昇順に入れ替えることになります。この操作で全体を昇順にできなければどうやっても昇順に入れ替えることはできません。

図で操作をイメージするとこの通りです。

実装では、K個のリストを作成し入れ替え可能なグループに分類します。その後これらをグループごとに昇順にソートしてから、一つずつ順番に取り出します。この時、前の値 ≦ 次の値 が常に成り立っていれば全体を昇順に入れ替えられると判定することができます。
また、グループごとに一つずつ取り出す処理は zip() で簡単に実装出来ます。ただし、この時どのグループも同じ長さである必要があるので、全体の長さが K の倍数になるように、入れ替えに影響しない適当な数を Aに追加しておきます。

計算量は $O( Klog(N//K) + N )$で十分高速です。

コード

pypy3

N,K=[int(nk) for nk in input().split()]
# 全体の長さを Kの倍数に調整
A=[int(a) for a in input().split()]+[10**10]*(K-N%K)
X=[[] for _ in range(K)]

# 入れ替え可能グループに分類
for i in range(N+(K-N%K)):
    X[i%K].append(A[i])

# 各グループごとに昇順にソート
for i in range(K):
    X[i]=sorted(X[i])
    
# 直前の値
prev=0
# 全体が昇順か判定
for column in zip(*X):
    for a in column:
        if prev>a:
            print("No")
            exit()

        prev=a

print("Yes")

補足

• zip
  Python, zip関数の使い方: 複数のリストの要素をまとめて取得

D.Together Square

問題

難易度： 緑色 1191

考察

i=1 から順番に、 条件を満たす j が何か考えてみます。

よって、$j= i\ ×$ 平方数 となりそうです。
この仮定の是非を検証します。

a , b を指数が奇数になる素因数の積、 x , y を指数が偶数になる素因数の積( 平方数 )として、$i=ax , j=by$ と置くことができます。この時 $ i × j = ab\ ×$ (平方数) となるので、$i × j$ が平方数となる条件は ab が平方数になること になります。
これは $a=b$ でのみ成立します。なぜなら、$a≠b$　の場合、必ず a または b のどちらかにのみ素因数 p が存在することになり、この素因数 p の指数は偶数にすることができないからです。
したがって $i = ax , j=ay\ $ となります。

$x =1$ の場合、 $j = i$ × 平方数　が成立するので、 これが 「 $i × j$ が平方数 」となるための必要十分条件であることを示せました。

ここで、この $i$ と $j$ とで表現できる $(i,j)$ の組の数は ( 平方数の数 )^2 個存在します。
実は $x≒1$ となる $i$ では、$j = i$ × 平方数 は必要十分条件にならないのですが、$x=1$ の時点でそのような $i$ からなる組も含めて数えておけば、$x=1$ の場合に限定して考えることができるようになります。

以上より、$i$ を全探索して、その $i$ が一度も $j$ として登場していなければ $N//i$ 以下の平方数の個数を求めることで、条件を満たす $(i,j)$ の組の総数を数えられるということが分かりました。

実装では、二分探索で　$N//i$ 以下の平方数の個数を求めることにしました。
この時、1 ～ N のすべてに対して ➀ 平方数二分探索 ➁ 探索済みチェック の処理が生じる可能性があるが、この計算量は $O( N logN )$ であるので十分高速です。

コード

pypy3

N=int(input())
ans=0
seen=[False for _ in range(N+1)]
for i in range(1,N+1):
    # 一度見た i はスキップ
    if seen[i]:
        continue

    # 二分探索で i × 平方数 ≦ N を満たす平方数の個数を求める
    ok=0
    ng=N+1
    while ng-ok!=1:
        mid=(ng+ok)//2
        j=i*mid**2
        if j <= N:
            ok=mid
        else:
            ng=mid
    
    ans+=ok**2
    
    # 平方数として登場した値は探索しないようにする
    for x in range(1,ok+1):
        seen[i*x**2]=True

print(ans)

別解

$ x ≒ 1$ の場合が面倒なので、あらかじめ 1 ～ N までのすべての数を平方数で割れるだけ割っておきます。
これによって、$i=j$ が条件を満たすための必要十分条件となります。
したがって、平方数で割り終えた後に同じ値となるものの個数で組の個数を求めることができます。

実装では、エラストテネスの篩風に、平方数で割るべき値を効率的に調べています。
平方数で割れるだけ割るための回数は、雑に見積もっても $log_4N$なので、全体として計算量は $O(NlogN × log_4N)$ となって十分高速です。

コード

pypy3

from collections import Counter
N=int(input())
table=[i for i in range(N+1)]
for j in range(2,N+1):
    # 一度平方数で割ったものはスキップ
    if table[j]!=j:
        continue
    for k in range(j**2,N+1,j**2):
        while table[k]%j**2==0:
            table[k]//=j**2

# 同じグループの個数をカウント
X=Counter(table[1:])
ans=0
# 同じグループの個数^2 個 (i,j) の組はある
for x in X.keys():
    ans+=X[x]**2

print(ans)

計算量詳細

$K = p_1^{a_1}\ ×\ p_2^{a_2}$ ... とする。 $(\ p_1 ,p_2$ は $\ p_1<p_2<... $ を満たす平方数)

両辺に対数を取ると

\begin{align}
logK &= a_1 × logp_1 + a_2 × logp_2 +...\\
&≧ (a_1+a_2+...) × logp_1\\
\end{align}

ここで底を変換すると

\begin{align}
log_{p1}K  = a_1 + a_2+...
\end{align}

ここで、右辺は平方数の個数を表していているので、1 ～ N までそれぞれ高々 $log_4{N}$ 回しか割れないことを示せました。

また、$2^2$ で割るべき数は N//4 個 , $3^2$ で割るべき数は N//9 個であるので、平方数で割る数の個数 $C$ は

\begin{align}
 \\\\[0pt]C\ &= \frac{N}{4} + \frac{N}{9} + \frac{N}{25}...\\
 \\\\[0pt]&≦ \frac{N}{2} + \frac{N}{3} + \frac{N}{4}...\\ 
 \\\\[0pt]&=NlogN
\end{align}

となります。したがって全体の計算量は最悪でも $O(NlogN × log_4N)$ と推定できました。

補足

• 二分探索
  ・【AtCoder】ABC271のA,B,C,D における Python解説
  こちらC問題 別解で詳細に解説しています。

  ・二分探索アルゴリズムを一般化 〜 めぐる式二分探索法のススメ 〜
  最もわかりやすく二分探索を解説している記事なので必読です。

  ・さまざまなアルゴリズム設計技法
  アルゴ式で練習してマスターしましょう。

• 個数管理
  PythonのCounterでリストの各要素の出現個数をカウント
  個数管理するときに便利なので覚えておくとよいと思います。

• 計算量推定
  ・調和級数1+1/2+1/3...が発散することの３通りの証明
  調和級数の部分和は $logN$ に近似します

  ・バーゼル問題の初等的な証明
  平方数の逆数和よりも調和級数が大きいことを利用して計算量を推定しましたが、実は平方数の逆和は 2以下であることが知られています。

• エラストテネスの篩
  素数判定が $O(loglogN)$ で行えるアルゴリズムです。
  エラトステネスの篩 (ふるい) を徹底解説 〜 実装から計算量まで 〜

E.Small d and k

問題
難易度 ： 水色 1202

考察

次数は頂点に連結している辺の数を表します。この問題では与えられるグラフの次数は 3 と非常に小さいです。また、考えるべき距離 $k_i$ も 3以下と非常に小さいです。
つまり、頂点 $x_i$ から距離 $k_i$ 以下の頂点は、下図で示すように高々 22個しかありません。

したがって、クエリごとに素直にこれらを全て調べても　計算量は $O(N)$ となって十分高速です。

特定の距離までの探索は BFS で実装すればよいでしょう
ただし、頂点間距離を表すリストをクエリごとに新しく作成してしまうと $O(N^2)$ になってしまう点に注意する必要があります。
このように全体をわざわざ初期化しなくても、そのクエリで探索した頂点の位置のみ初期化するだけで十分です。この工夫を施せば制限時間に間に合う実装が可能です。

pypy3

from collections import deque
N,M=[int(nm) for nm in input().split()]
G=[[] for _ in range(N)]
for _ in range(M):
    a,b=[int(ab)-1 for ab in input().split()]
    G[a].append(b)
    G[b].append(a)

Q=int(input())
dist=[-1 for _ in range(N)]
for _ in range(Q):
    x,k=[int(xk) for xk in input().split()]
    que=deque()
    que.append(x-1)
    ans=x
    # 探索した頂点を覚えておく
    seen=set()

    dist[x-1]=0
    seen.add(x-1)
    while que:
        now=que.popleft()
        if dist[now]==k:
            continue
        for nex in G[now]:
            if dist[nex]<0:
                que.append(nex)
                dist[nex]=dist[now]+1
                ans+=nex+1
                seen.add(nex)
    
    
    print(ans)
    # 探索した頂点の位置の距離のみ初期化
    for i in seen:
        dist[i]=-1

補足

なし

終わり

2022年、もう一本記事だします !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!