日能研の問題をPythonで解いてみた(2018/1)

nichinoken.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-

import itertools
from functools import reduce
from operator import add
from fabric.colors import red


#上下左右のマスの取得
def vectol_list(list_id, size):
	v_work = []
	v_id_x,v_id_y = (list_id%size, int(list_id/size))
	for v_x,v_y in ((0,-1),(0,1),(-1,0),(1,0)):
		v_work_x = v_id_x + v_x
		v_work_y = v_id_y + v_y
		if (-1 < v_work_x < size) and (-1 < v_work_y < size):
			v_work.append(v_work_y*size + v_work_x)
	return v_work


#結果の出力
def anser_print(board, anser):
	size = len(board[0])
	for ans_id,ans_val in enumerate(anser):
		print("-"*10)
		print("No:{0} {1}".format(ans_id+1, ans_val))
		board_work = [i for i in "".join(board)] # リストに変換する
		for ans in ans_val:
			board_work[ans] = red(board_work[ans]) # 赤色に変換する
		for x in range(size, size*size, size+1): 
			board_work.insert(x,"\n") # 改行コードを追加する
		print("".join(board_work))


#実装方法1
def method1(board, set_num, sum):
	anser = []
	size = len(board[0])

	# 手順１
	work_dic = {"0":[], "1":[]}
	for id,val in enumerate("".join(board)):
		work_dic[val].append(id)

	# 手順2,手順6
	for check in itertools.combinations(work_dic["0"], set_num):
		# 手順3
		black_list = set(work_dic["1"]) | set(check)
		# 手順４,手順5
		p = reduce(add, list(map(lambda a:len(set(vectol_list(a,size)) & black_list), black_list)))
		if sum == p:
			anser.append(check)
			anser.append(check)

	# 手順7
	anser_print(board, anser)


#実装方法2
def method2(board, set_num, sum):
	anser = []
	size = len(board[0])

	# 手順1
	work_dic = {"0":[], "1":[]}
	for id,val in enumerate("".join(board)):
		work_dic[val].append(id)

	# 手順2
	temp = set(work_dic["1"])
	pre_p = reduce(add, list(map(lambda a:len(set(vectol_list(a,size)) & temp), temp)))

	# 手順3
	default_put = list(map(lambda a:vectol_list(a,size), work_dic["1"]))
	default_put = list(set(reduce(add, default_put)) - set(work_dic["1"]))

	# 手順4
	check_list1 = [list(i) for i  in itertools.combinations(default_put, set_num)]

	# 手順5
	check_list2 = []
	temp = list(map(lambda a:vectol_list(a,size), default_put))
	temp = list(map(lambda a:set(a) - set(work_dic["1"] + default_put), temp))
	temp = list(map(lambda a:list(set(a) - set(work_dic["1"] + default_put)), temp))
	for id, val in  enumerate(default_put):
		for c in temp[id]:
			check_list2.append([val,c])

	# 手順6, 手順9
	for check in check_list1 + check_list2:
		# 手順7
		temp_p = pre_p
		black_list = work_dic["1"][:]
		for c in check:
			black_list.append(c)
			temp_p += len(set(vectol_list(c,size)) & set(black_list)) * 2
		# 手順8
		if temp_p == sum:
			anser.append(check)

	anser_print(board, anser)


def main():
	board = [
		"000000",
		"001000",
		"011010",
		"001110",
		"001010",
		"000000",
	]
	set_num = 2
	output_condition = 26

	method1(board, set_num, output_condition)
	method2(board, set_num, output_condition)


main()