任意精度平方根と素数テスト (SRM821 CrossPrimesOut)

独特な要素が多く、結構印象深い問題。出題されたのは１年前。
出題はSRM 821 Div1の3問目。

問題

community.topcoder.com
正整数Xが1つ与えられる。その平方根を取って小数点を無視して、無限の長さの数字の文字列が得られる。

例 : sqrt(47) → 685565460040104...

次に、無限に存在する素数に対して小さい順に、以下の手続きを行う。

その素数が文字列の中で最初に現れる位置を見つける
その文字を空白で置き換える。（無限に長い文字列の中にその素数が存在しない場合は何もしない）

この無限の手続きを終えると空白区切りのトークン列が得られる。各トークンは数字の文字列になっている。

平方数ではない正整数Xと整数kが与えられるので、k番目のトークンを求めよ。

制約

1 ≦ X ≦ 10000
X は平方数ではない
0 ≦ k ≦ 30

例

X = 47
k = 1

答 5654600

6855654600401044124935871449084848960460643461001326275485108185678...
68_5654600___04___49_58_1___08484__604606__4__00__2627_____08185__8...

（0オリジンで）k 番目のトークンは 5654600。
同じ素数は最初の一回しか消さない。一度空白に置き換えたい値は別の素数にはならない。

解法

400桁の平方根を求めるて、十進数で文字列とする。
1000000 以下の素数について、文字列に存在するなら空白に置き換える。
各トークンが7文字の素数を含むなら空白に置き換える。以下8文字、9文字、と順に調べる。

公式解説
docs.google.com

問題の独自性

任意精度小数を使うというのが珍しい。ただしsqrtを求めるだけなのでニュートン法や二分法だけできればいい。
この問題の制約の範囲では400桁でいいが、実行してみるまで何桁必要か予想できない。
無限の文字列から見つからないときは何もしないが、見つかるか見つからないか不明な時はどうするのか？答えには影響しないので無視すればいいが不安定。
もし100文字以上のトークンが残ったらこの中に素数がないことを判定しなければならない。実装中はトークンの長さが予想できない。

実装の注意

任意精度小数のある言語を選択するのが良い。
結構大きめの素数を判定する必要があるので多倍長を用意する。(例えば7936145404496840003051、2^64 より大きい素数が出てくる)
文字列から素数を探すので誤って "0000002" を7桁の素数と判定しないように。

コード (python3)

from decimal import *

getcontext().prec = 500

MAX = 1000000

prime = []
tbl = [True] * MAX
tbl[0] = False
tbl[1] = False
for i in range(2, MAX):
    if tbl[i]:
        prime.append(i)
        for j in range(i*i, MAX, i):
            tbl[j] = False


def isSprp(a, s, d, n):
    a %= n
    if a == 0:
        return True

    x = pow(a, d, n)

    if x == 1:
        return True

    for i in range(s):
        if x == n - 1:
            return True
        x = x * x % n

    return False


def isPrime(n):
    if n < MAX:
        return tbl[n]

    if n % 2 == 0 or n % 3 == 0 or n % 5 == 0 or n % 7 == 0:
        return False

    d = n - 1
    s = 0
    while d % 2 == 0:
        s += 1
        d //= 2
    
    bases = [2, 325, 9375, 28178, 450775, 9780504, 1795265022]
    for b in bases:
        if not isSprp(b, s, d, n):
            return False

    return True


class CrossPrimesOut:
    def findTerm(self, X, N):
        s = str(Decimal(X).sqrt()).replace('.', '')

        for p in prime:
            p = str(p)
            s = s.replace(p, ' ', 1)

        s = s.split()

        if not s:
            return []

        maxLen = max([len(e) for e in s])

        for curLen in range(7, maxLen+1):
            while True:
                cand = []
                for i in range(len(s)):
                    if len(s[i]) < curLen:
                        continue
                    for j in range(len(s[i]) - curLen + 1):
                        if s[i][j] == '0':
                            continue
                        num = int(s[i][j:j+curLen])

                        if isPrime(num):
                            cand.append((num, i, j))

                if not cand:
                    break

                num, i, j = min(cand)

                if len(s[i]) == curLen:
                    s.pop(i)
                elif j == 0:
                    s[i] = s[i][curLen:]
                elif len(s[i]) == j + curLen:
                    s[i] = s[i][:j]
                else:
                    left = s[i][:j]
                    right = s[i][j+curLen:]
                    s[i] = right
                    s.insert(i, left)

        return int(s[N]) % 1000000007