CodeIQ：最寄りの素数

問題の概要

私のプログラム

Pythonで解答しています。

#!/usr/local/bin/python3
# -*- coding:utf-8 -*-

import sys
import math
import itertools
import fileinput

# 値から並びの最初の値を得る
def getColFirst(n):
	return n*n+1

# 値から並びの最後の値を得る
def getRowLast(n):
	return n*n+2*n+1

# 値から現在の座標を得る
def getRowCol(n):
	sq = math.ceil(math.sqrt(n))
	for i in range(sq):
		first = getColFirst(i)
		last = getRowLast(i)
		mid = (first + last) // 2
		if first <= n <= last:
			break

	if n == mid:
		return [i, i]
	elif n < mid:
		return [i-(mid-n), i]
	else:
		return [i, i-(n-mid)]

# 座標の値を得る
# @param pos 座標の配列[row, col]
def getValueOfPos(pos):
	i = max(pos)			# 現在の並び
	if pos[1] >= pos[0]:	# 縦の並びにある場合
		return getColFirst(i) + min(pos)
	else:
		return getColFirst(i) + (2*max(pos)-min(pos))

# 指定した座標（pos）からdだけはなれた場所の値のリストを返す
# @param pos 座標の配列[row, col]
# @param d 何マスはなれているところを調べるか
def getValueList(pos, d):
	st_r = pos[0] - d if (pos[0]-d) >= 0 else 0
	ed_r = pos[0] + d
	st_c = pos[1] - d if (pos[1]-d) >= 0 else 0
	ed_c = pos[1] + d

#	print("st_r=" + str(st_r) + " ed_r=" + str(ed_r) +", stc=" + str(st_c) + " ed_c=" + str(ed_c))

	ret = []
	for i in range(st_r, ed_r+1):
		if(isCorrectDistance(pos, [i,st_c], d)):
			ret.append(getValueOfPos([i,st_c]))
		if(isCorrectDistance(pos, [i,ed_c], d)):
			ret.append(getValueOfPos([i,ed_c]))
	for j in range(st_c+1, ed_c):
		if(isCorrectDistance(pos, [st_r,j], d)):
			ret.append(getValueOfPos([st_r,j]))
		if(isCorrectDistance(pos, [ed_r,j], d)):
			ret.append(getValueOfPos([ed_r,j]))

#	print("周囲:" + str(sorted(ret)))
	return ret

# targetが正しくbaseからdだけはなれたところにあるかチェックする
def isCorrectDistance(base, target, d):
	if (math.fabs(base[0]-target[0]) != d) and (math.fabs(base[1]-target[1]) != d):
		return False
	return True

# 素数かどうかを判定する
def isPrime(q):
	if q == 2: return True
	if q < 2 or q&1 == 0: return False
	return pow(2, q-1, q) == 1

# 距離の2乗を計算する
def squareDistance(center, pos):
	rd = center[0]-pos[0]	# 縦方向
	rc = center[1]-pos[1]	# 横方向
#	print("center:" + str(center) + " pos:" + str(pos) + " rd:" + str(rd) + " rc:" + str(rc))
	return rd*rd + rc*rc

# 現在のリストの中でposに最も近いもののリストを返す
# @param pos 中心座標
# @param lst 距離計測対象の座標リスト
# @return [距離, [素数の配列]]
def getNearList(pos, lst):
	if len(lst) == 0:
		return [2 ** 63 - 1, []]

	dist = []
	for i in lst:
		dist.append(squareDistance(pos, getRowCol(i)))

#	print("素数" + str(lst))
#	print("距離" + str(dist))

	# 距離の2乗の中で最も近い物だけを抽出する
	ret = []
	m = min(dist)
	for i,k in enumerate(dist):
		if k == m:
			ret.append(lst[i])

	return [m,ret]

# 最も近傍の素数を探す
def getNearlyPrimes(n):
	pos = getRowCol(n)	# 行、列の位置を求める

	# 1マスずつ拡大しながら周辺を探索して最も近い場所で見つかったものを返す
	ret = []
	i=1				# 中心から何マス離れるか
	m=2 ** 63 - 1	# 距離の最小値（ループ終了チェックに使う）
	while True:
		# 最も近いマスまでの距離が結果として得られたものの距離よりも大きくなったらやめる
		if (len(ret) != 0) and (i*i > m):
			break

		v = getValueList(pos, i)
		i+=1

		l = list(filter(isPrime, v))	# 1マスずつ広げながら素数を探す
		d, ln = getNearList(pos, l)		# 見つかった中で最も近傍のものを抽出する

		if len(ln) == 0:	#抽出したリストが空なら次
			continue
		elif len(ret) == 0:	#結果が空なら初めての要素なのでチェックなしに追加
			ret += ln
			m = d
		else:
			if m == d:	# 新規取得した要素が現在と同じ距離にあるなら追加する
				ret += ln
			elif d < m:	# 新規取得した要素の方が現在の結果より小さい場合、結果を作り直す
				ret = []
				ret += ln
				m = d

	return sorted(ret)

# 結果出力
def printResult(ret):
	s = ""
	for i, k in enumerate(ret):
		s += str(k)
		if i != len(ret)-1:
			s += ","
	print(s)

#==============================================================================
# main
#------------------------------------------------------------------------------
if __name__ == "__main__":
	for line in fileinput.input():
		if not line.strip():
			continue

		n = int(line)
		ret = getNearlyPrimes(n)
		printResult(ret)

解説

この問題はかなり難しいと思います。しかし、問題はわかりやすいですし、数学的にもそんなに難しくはありません。

雑感

かなり難しい問題ではあったのですが、やればなんとかなると言う感じだったのでなかなか楽しめました。
しかし、素数判定が完全ではないことに気づいたのは別の時に素数判定のルーチンを流用しようとした時でした。その時には調べ直して冷や汗をかきました。