EAFP

~ Easier to Ask for Forgiveness than Permission ~

Atcoder ABC284をpythonで解いてみた

目的

Atcoder ABC284に出場したので、その結果、感想、復習を兼ねて記事を書きました。

特に、緑になってから約10ヶ月程度立ちましたが、rateが上下し一向に向上の兆しが見られない原因として 復習が足りないのではないか?という仮説を立て、今年からはrateが下がった際には、記事を書く!というのを自分に課していこうと思っています。

atcoder.jp

解いたコードで、気になって色々修正もしたものは下記に置いてあります。動作確認済みです。

github.com

結果

ABCの3完で、D問題は方針は概ね合っていたものの、考察が足りないため解けず(後述します) E問題に移ったが、時間内に解くことができずタイムアップとなってしまいました。

図1. ABC284の結果

ABCまでは、10分程度で解け、Dも方針立てるのに5分程度でしたし、E問題も緑diffであったことを考えると、5完したかったところですが、これができないあたりに課題を感じました。

自分の今のレベルからすると、A, B問題は簡単であったと感じるので、解答のみに留めます。

  • A問題

atcoder.jp

  • B問題

atcoder.jp

C問題

N 頂点 M 辺の単純無向グラフのうちに、連結成分が何個あるか?という問題

atcoder.jp

考察

連結成分とは何か?に関しては、下記に記載があるので、ここでの説明は割愛させていただきます。

algo-method.com

まず、連結成分を求める際に Union-Find を使おうと考えました。Union-Findを利用することで、辺で繋がれた頂点同士を1つのグループとして扱うことができるためです。 ここでは、DFSや、BFSを使って要素を探索することで、解答することでも可能だと考えましたが

amateur-engineer-blog.com

Union-Findを使えばより簡単(手元にUnion-Findを実装したクラスがあり、与えられた辺を近接リストにして処理せずとも、辺の情報のみで解答可能)にできると考え、下記のように実装しました。

# いつも使っているUnionFindのClass(下記を参照にし、一部変更したりして利用している)
# https://note.nkmk.me/python-union-find/
class UnionFind():
    def __init__(self, n):
        self.parents = [-1] * n

    def find(self, x):
        if self.parents[x] < 0:
            return x
        else:
            self.parents[x] = self.find(self.parents[x])
            return self.parents[x]

    def union(self, x, y):
        x = self.find(x)
        y = self.find(y)
        if x == y:
            return
        if x < y:
            x, y = y, x
        self.parents[x] += self.parents[y]
        self.parents[y] = x


N, M = map(int, input().split())
uf = UnionFind(N)
for _ in range(M):
    u, v = map(int, input().split())
    # 要素同士を1つにまとめる処理
    uf.union(u - 1, v - 1)
# N個の要素の親を確認し、親要素の種類を数え上げる
ans = set()
for i in range(N):
    p = uf.find(i)
    if p not in ans:
        ans.add(p)
# 親の要素の種類数が答えになるので、出力する
print(len(ans))

atcoder.jp

これにより、時間内に問題を解くことができました。 時間計算量: O(N + M)、空間計算量 O(N + ansに格納される要素数) であると考えます。

実際は、Union-Findのfindを繰り返して、計算するよりも、Union-Findのクラス内で保持している self.parents のうち、-1 である要素を数え上げることでも、答えはわかる上に、都度findしない分早くなると思われます。 しかし、find自体は、下記の両方の工夫をして実装しており、実質 O(1) 程度の速度で実行可能であると考えると、ほぼ差は出ないと考えました。

algo-method.com

が、後々やってみる(下記がカイゼン後の提出結果)と50ms程度の速度向上が見られました。 これは、range(N) を作成したり、ans 内に すでに同じ要素があるか? findでの処理はO(1)に近づくとはいえ、完全に1ではない といったあたりが関係していると思われます。

atcoder.jp

D問題

正整数 N =p2 × q (素数 p,q) と因数分解可能な N が T個与えられるので、それぞれで、p, q を求めて出力せよ!という問題

atcoder.jp

考察

まず、 1 <= N <= 9 × 10^18 と非常に巨大であることから、N全てを探索することも、ルートを取ったとしても間に合わないと思いました。

つまり、通常よく実施する素因数分解では間に合わないため、別の方法を考える必要があります。

amateur-engineer-blog.com

ちなみに、解説にもある ポラード・ロー素因数分解法という高速に素因数分解する方法を使えば解けるようですが、自分は知らなかったためそのような選択はしませんでした。

次に着目した点に関しては、1/3 乗 であれば 106 程度に収まるため、計算可能ではないか?と思い、エラストてネスの篩を用いて、N1/3 の範囲の素数を求め、それらを使って素因数分解できないか?を考えました。

club.informatix.co.jp

理由としては、 N = p * p * q となるため、p か q のどちらかは必ず上記の範囲内の素数として発見できると考えたためです。 ここまでは良かったのですが、問題を解いている際は、pにのみ着目して考えてしまい、pが小さい際には、qはこの範囲を超えてしまう・・・どうしよう・・・と思って色々考えて時間を潰してしまいました。

反省としては、p か q のどちらかになるのであれば、その逆を求めることは簡単なはずなのに、その発想に至らず解くことができなかった。後々言われるとわかるけど、その場では思いつかなかったと言うところになります。

では、どうすればいいのか?というと、106 程度 の素数のうち、Nを割り切ることができる最大の数を x としておいたときに - N % x2 = 0 であれば、p = x である - そうでなければ、q = x であり、 p = √(N/x) である

ということになり、問題を解くことができます。

from math import sqrt

# エラストてネスの篩を用いて素数を列挙する関数
def sorted_prime_list(limit):
    limit += 1
    primes, is_prime = [], [True] * limit
    for i in range(2 ** 2, limit, 2):
        is_prime[i] = False
    for p in range(3, limit, 2):
        if not is_prime[p]:
            continue
        primes.append(p)
        for i in range(p * p, limit, p):
            is_prime[i] = False
    return primes


T = int(input())
N = [int(input()) for _ in range(T)]
# 何度も素数列挙をするのは効率が悪いので、1度で実施する
primes = sorted_prime_list(min(max(N), 3 * pow(10, 6)))
for n in N:
    p_or_q = 2
    for prime in primes:
        # prime が n^(1/3) を超えてしまうような範囲は探索する必要がないのでbreakする
        if pow(n, 1 / 3) + 1 < prime:
            break
        if n % prime == 0:
            p_or_q = max(p_or_q, prime)
    ans = (p_or_q, n // p_or_q ** 2) if n % p_or_q ** 2 == 0 else (round(sqrt(n // p_or_q)), p_or_q)
    print(*ans)

atcoder.jp

これにより、問題を解くことができました。 時間計算量: O(XloglogX + T * primesのlength * round(sqrt)の計算量 )、空間計算量 O(x + x以下の素数の数 + 2 * T) であると考えます。 X = min(max(N), 3 * pow(10, 6)) とし、sqrtの計算量に関しては色々調べましたが log程度かそれ以下なのはわかったのですが、わからず・・・

ちなみに、高速で平方根も止める方法は色々あるようでその辺りは興味深かったです。

takashiijiri.com

E問題

N 頂点 M 辺の単純無向グラフを探索し、パスの本数を数えよ!という問題

atcoder.jp

考察

まず、要素1 からの経路を探索していくことを考えたので DFS を用いて解こうと考えました。pathの数をKとした際に、求める答えが min(K,10 ^6) であるが故に、106 以上のpathを回る必要がないということが条件に書かれているためです。

そのため、書き慣れた再帰を用いたDFSの関数を書き、通った要素を再度通らないように vis というリストを用意し、通った箇所は True とするようにして状態管理を行いました。

from collections import defaultdict
from sys import setrecursionlimit

from pypyjit import set_param

set_param("max_unroll_recursion=-1")
setrecursionlimit(pow(10, 9))


def dfs(cur):
    global ans
    if ans >= limit:
        print(limit)
        exit(0)
    for nxt in edge[cur]:
        if not vis[nxt]:
            vis[nxt] = True
            ans += 1
            dfs(nxt)
    vis[cur] = False


N, M = map(int, input().split())
edge = [[] for _ in range(N)]
for _ in range(M):
    u, v = map(int, input().split())
    edge[u - 1].append(v - 1)
    edge[v - 1].append(u - 1)

ans, limit = 1, pow(10, 6)
vis = defaultdict(bool)
vis[0] = True
dfs(0)
print(ans)

しかし、何度やってもTLEしてしまいます・・・

atcoder.jp

ダメ元で、python3で提出したら通りました。

atcoder.jp

pypyでの再帰は遅い問題は指摘されていましたが、pypyjit でパラメータを設定することで回避できるという話も見かけていて、今までこれを使って回避できたと思っていたのですが

qiita.com

今回に限ってなのか、そもそもさほど効果がなかったのか?の深掘りはできていませんが、こういうこともあるのだろうな。というところで、調査はやめてしまいました。

doc.pypy.org

もし、この辺り詳しい方いれば教えて欲しいです。

まとめ

今回は、Dでコケてしまい、その後軌道修正できずに終わってしまうという悔しい結果になりました。整数問題というか、Dのようなタイプの問題が苦手て、自分のrateよりも低い問題でも解けなかったりすることが多いので、この辺り解けるようにしていければと思っています。

また、考察をより深く!早く!立てて手戻りなく、スムーズに解けるようになるというのも課題なので、コンテスト中も、普段別の問題を解く時も、復習で記事を書く際も、この辺りの能力を高めていけるように、頑張っていければと思いました。

いつか、水色になれる日を夢見て、今後もAtcoderを頑張っていければと思います。