ABC475Dを解いた【エラトステネスの篩・DFS全探索】

筆者はレート800前後の茶～緑コーダ

ABC475の問題を解いていく

実装コード

文字への数字の割り当てをDFSで探索し、エラトステネスの篩で素数かどうかを判定する。

• 文字列 S を受け取り、各文字にIDを付けて配列 C に保存する
• sieve(10 ** N - 1) で、N 桁以下の整数の素数判定表 prime を作る
• digit で各文字に割り当てた数字、used で使用済みの数字を管理する
• dfs(i, now) で左から i 文字分を数字に置き換えた値 now を持ち、割り当て済みの文字なら同じ数字を使って次へ進む
• 未割り当ての文字には未使用の数字を小さい順に試し、先頭が 0 になる場合は除外する
• 再帰先で素数が見つかればその値を返し、見つからなければ digit と used を元に戻して次の数字を試す
• 全文字を置き換えたら prime[now] で判定し、最後に見つかった最小の素数を出力する。見つからなければ -1 を出力する

main.py

from bisect import bisect_left, bisect_right, insort_left, insort_right
from collections import defaultdict, Counter, deque
from functools import reduce, lru_cache
from itertools import product, accumulate, groupby, combinations
import sys
import os
def rI(): return int(sys.stdin.readline().rstrip())
def rLI(): return list(map(int,sys.stdin.readline().rstrip().split()))
def rI1(): return (int(sys.stdin.readline().rstrip())-1)
def rLI1(): return list(map(lambda a:int(a)-1,sys.stdin.readline().rstrip().split()))
def rS(): return sys.stdin.readline().rstrip()
def rLS(): return list(sys.stdin.readline().rstrip().split())
IS_LOCAL = int(os.getenv("ATCODER", "0"))==0
err = (lambda *args, **kwargs: print(*args, **kwargs, file=sys.stderr)) if IS_LOCAL else (lambda *args, **kwargs: None)

def sieve(n):
    is_prime = [True for _ in range(n+1)]
    is_prime[0] = False
    is_prime[1] = False

    for i in range(2, int(n**0.5)+2):
        if is_prime[i]:
            j = 2 * i
            while j <= n:
                is_prime[j] = False
                j += i
    table = [ i for i in range(1, n+1) if is_prime[i]]
    return is_prime, table

def main():
    S = rS()
    N = len(S)

    dic = {}
    C = []
    for s in S:
        if s not in dic:
            dic[s] = len(dic)
        C.append(dic[s])

    prime, _ = sieve(10 ** N - 1)

    digit = [-1] * len(dic)
    used = [False] * 10

    def dfs(i, now):
        if i == N:
            return now if prime[now] else -1

        c = C[i]
        if digit[c] != -1:
            return dfs(i + 1, now * 10 + digit[c])

        for d in range(10):
            if used[d] or (i == 0 and d == 0):
                continue
            digit[c] = d
            used[d] = True
            ans = dfs(i + 1, now * 10 + d)
            if ans != -1:
                return ans
            digit[c] = -1
            used[d] = False
        return -1

    print(dfs(0, 0))
if __name__ == '__main__':
    main()

感想

文字ごとの割り当てと使用済みの数字を分けて管理すると、
同じ文字には同じ数字、違う文字には違う数字という条件を扱いやすい。
素数判定を篩で前計算しておけば、DFSでは数字の割り当てに集中できるのがポイントだった。