logic_puzzle_power.py

# [2017-01-05] 基準となるプログラム
# [2017-01-06] 全列挙探索(左詰め)を実装
import subprocess
import matplotlib.pyplot as plt
import matplotlib.ticker as ticker
from time import sleep
from sys import argv
from itertools import product

# Global
R = None
C = None
row_keys = None
col_keys = None


def usage():
    print('python3 {f} puzzle_file'.format(f=__file__))


def main(puzzle_file):
    # 問題の入力
    get_puzzle(puzzle_file)

    # rowinfo[keyix] = (r, ix)
    rowinfo, num_keys = get_rowinfo()

    # 全列挙
    B = power(rowinfo, num_keys)

    printB(B, 'Finish!')


def power(rowinfo, num_keys):
    imgix = 0
    # key2pos[keyix] = (r, c, key)
    key2pos = set_B(0, rowinfo, num_keys) # 初期状態
    B = gene_B(key2pos, num_keys)
    while not feasible(B):
        printB(B, imgix); imgix += 1
        l = 0
        keyix = num_keys - 1
        while not can_slide(keyix-l, rowinfo, key2pos):
            l = l + 1
        key2pos[keyix-l][1] += 1 # keyix-lを1つ右にスライド
        key2pos = set_B(keyix-l+1, rowinfo, num_keys, key2pos)
        B = gene_B(key2pos, num_keys)
    return B


def gene_B(key2pos, num_keys):
    B = [[-1 for c in range(C)] for r in range(R)]
    for keyix in range(num_keys):
        r, c, key = key2pos[keyix]
        for k in range(key):
            B[r][c+k] = 1
    return B


def set_B(keyix, rowinfo, num_keys, key2pos=dict()):
    # update key2pos
    if keyix < num_keys:
        for kix in range(keyix, num_keys):
            r, ix = rowinfo[kix]
            cor_key = row_keys[r][ix]
            if ix == 0:
                key2pos[kix] = [r, 0, cor_key]
            else:
                pre_r, pre_ix, pre_key = key2pos[kix-1]
                key2pos[kix] = [pre_r, pre_ix+pre_key+1, cor_key]
    return key2pos


def can_slide(keyix, rowinfo, key2pos):
    r, ix = rowinfo[keyix]
    _, c, _ = key2pos[keyix]
    keys = row_keys[r]
    if c < C-sum(keys[ix:])-len(keys[ix:])+1:
        return True
    else:
        return False


def get_rowinfo():
    dic = dict()
    keyix = 0
    for r in range(R):
        keys = row_keys[r]
        for ix in range(len(keys)):
            dic[keyix] = (r, ix)
            keyix = keyix + 1
    return dic, keyix


def feasible(B):
    """Bが列に関して, 制約を満たしているかを調べる"""
    for c in range(C):
        e = ''
        for r in range(R):
            e += ['a', 'b'][B[r][c] == -1]
        e = [elm for elm in e.split('b') if elm]
        e = list(map(len, e))
        if col_keys[c] != e:
            return False
    return True


def printB(B, msg=''):
    subprocess.run(['clear'])
    print('\n'*5)
    print(msg)
    for r in range(R):
        for c in range(C):
            if B[r][c] == 1:
                sq = '■'
            elif B[r][c] == -1:
                sq = ' '
            else:
                sq = '-'
            print(sq+' ', end='')
        print()
    sleep(0.05)


def imageB(B, title=None, savename=None):
    fig = plt.figure(figsize=(6, 6))
    ax = fig.add_subplot(111)
    ax.set_xlim([0, C])
    ax.set_ylim([0, R])
    ax.xaxis.set_major_locator(ticker.NullLocator())
    ax.yaxis.set_major_locator(ticker.NullLocator())

    for r, c in product(range(R), range(C)):
        if B[r][c] == 1:
            pos = (c, R-r-1)
            ax.add_patch(plt.Rectangle(pos, 1, 1, color='k'))
        if B[r][c] == 0:
            x = (c+0.3, c+0.7)
            y = (R-r-0.5, R-r-0.5)
            ax.plot(x, y, color='k', linewidth=0.2)

    if title is not None:
        ax.set_title(title)
    if savename is not None:
        plt.savefig(savename)
    else:
        plt.show()

    plt.close()


def get_puzzle(puzzle_file):
    global R, C, row_keys, col_keys
    row_keys = []
    col_keys = []
    for line in open(puzzle_file, 'r'):
        if not line.strip():
            continue
        sigh, *keys = line.split()
        if sigh == 'R':
            R = int(keys[0])
        if sigh == 'C':
            C = int(keys[0])
        if sigh == 'r':
            keys = list(map(int, keys))
            row_keys.append(keys)
        if sigh == 'c':
            keys = list(map(int, keys))
            col_keys.append(keys)


if __name__ == '__main__':
    if len(argv) == 2:
        main(argv[1])
    else:
        usage()