yukicoder : No.308 素数は通れません

Advent Calendar Contest初日から良い感じの問題。
http://yukicoder.me/problems/840

問題

1～Nの番号が書かれたマスを、左上から順に1段あたりW個で並べていくことを考える。
すなわち1段目は左から1～W、2段目は(W+1)～(2*W)…となっていく。

ここで、1番のマスから、隣接する合成数のマスを辿ってN番のマスに到達したい。
合成数Nが与えられるとき、条件を満たす最小のWを求めよ。

解法

1→2には移動できないので、初手は1→(W+1)に移動するしかない。
よってWの候補は(W+1)が合成数となるものである。

W=3,5,7は条件を満たすが、実際に数字を並べてみると到達可能なNはかなり限られていることがわかる。
Wを偶数に出来ると、偶数列目の縦のラインが(最初の2を除き)自由に縦移動できるようになる。
また、ある程度大きなNだと、この縦のライン同士を結ぶ奇数の合成数が登場するのでライン同士も移動可能になり、ほとんどの数字に移動可能になる。
Wが偶数で(W+1)が合成数となる最小のWは8である。

W=8のとき、Nが27以上なら縦の偶数ラインに任意に移動可能になる。
よって、

• Nが27未満ならWを総当たりしてみると良い。
• Nが27以上ならW=8が解の候補となりうる

ただし1つ落とし穴があって、N%8==1のとき、一番下の段はNのマスしかないので偶数列目の縦移動を駆使しても到達できない。
その場合、(N-8)が合成数である必要がある。

(N-8)が合成数でない場合、他に下の条件を満たす最小のWが解となる。

• Wが偶数
• (W+1)が合成数
• (N-W)が合成数

問題はNが大きく簡単に合成数判定できないことである。
そこでミラーラビン判定法を使うと良い。
合成数判定さえできれば、上記条件を満たすWを順に探して行けば良い。
(なお、writer解説によるとW=8とW=14のどちらかは必ず条件を満たすようだ。以下のコードはW=8以降愚直に探している)

import sys
import math
import random

N=input()

def create_erato(n):
	div = [0]*(n+2)
	primes = []
	for i in range(2,n+1):
		if div[i] == 0:
			primes.append(i)
			for j in range(i,n+1,i):
				div[j] = i
	return (div,primes)

div = create_erato(105)[0]
vis=[0]*105

def ok(N,M):
	q=[1]
	while len(q) > 0:
		v = q[-1]
		q.pop()
		
		for t in (v-1,v+1,v-M,v+M):
			if v%M==1 and t==v-1:
				continue
			if v%M==0 and t==v+1:
				continue
			if t<=0 or t>N:
				continue
			if div[t]!=t and vis[t]<M:
				vis[t]=M
				q.append(t)
	
	if vis[N]==M:
		return 1
	return 0

def MillarRabin(v,loop):
	d = v-1
	s = 0
	while d%2==0:
		s += 1
		d /= 2
	
	for i in range(loop):
		a = random.randrange(0,v-2)+1
		r = pow(a,d,v)
		if r==1 or r==v-1:
			continue
		ng = 0
		for j in range(s):
			r = pow(r,2,v)
			if r == v-1:
				ng = 1
		if ng==0:
			return 0
	
	return 1


if N>=100:
	random.seed()
	for v in range(8,100,2):
		if div[v+1]!=v+1 and (N%v!=1 or MillarRabin(N-v,100)==0):
			print v
			break
else:
	for M in range(2,105):
		if ok(N,M):
			print M
			break

まとめ

本番幸いミラーラビン判定法を思いつけて良かった。
Pythonのライブラリ持ってなかったので慌ててC++から移植した。