Google Code Jam 向け python template

明日は Google Code Jam Japan 2011 本戦です。
予選はなんとか突破できたので、Tシャツ目指してもうひと頑張りしてみます。

予選の前日にpython 向けテンプレートを公開しましたが、いざ予選で使ってみると、若干不満が残ったので修正してみました。修正箇所は以下の通りです。

1. 便利なライブラリを幾つか import
2. "-m" でマルチプロセス処理時に Ctrl + C で止められるようにした
3. "-q" を付けると、print 文による画面出力を抑制できるようにした
4. solve_xxx のような名前で複数の solver を実装しておいて、"-c"を付けて実行すると、それぞれの solver の出力を比較して相違が無いかを確認してくれるようにした

最後の奴は、「とりあえず small を力技で解く」とかやってしまった場合用です。
この場合、作りなおすなりリファクタリングするなりして、Large に耐える解答を作らなければならないわけですが、smallを正解できた solver を solver_small とでも rename しておけば、作りなおした関数が間違っていないかは -c を付けるだけで確認できるわけです。

まあ、いろいろ小細工を足してみましたが、最初から素直に正しい解答が書ける人なら無用の長物ばかりですね。orz

#!/usr/bin/env python2.7

import getopt
from collections import Counter, namedtuple
from itertools import *
import operator
import multiprocessing
import os.path
from random import *
import sys
from copy import deepcopy

# utils
def read_numbers():
    return map(int, infile.readline().split())
def read_string():
    return infile.readline().rstrip('\r\n')
def list2str(*array):
    return ' '.join(map(str,array))

############################################################

def read_precond():
    # NoTC is number of testcases.  others are preconditions.
    (NoTC,) = read_numbers()
    return locals()

# testcase parser
def read_testcase():
    # X,Y = read_numbers()
    # Ls = read_string()
    # Coff = namedtuple("Coff", ["x_id", "count", "limit", "satisfy"])
    #  for i in xrange(N):
    #    coffs.append(Coff(*([i,] + read_numbers())))
    # del i,x, y
    return locals()

# solver
def solve(Ls,**rest):
    # TODO: modify args to receive input data

    # idioms:
    #  Ls = [a*a for x in range(3,10,2)]
    #  map(lambda a,b: a*b, Ls, Ts)
    #  filter(lambda a: a!=1, Ls)
    #  product(range(1,X+1), range(1,Y+1))
    #  reduce(lambda s,a: s*2+a, range(3,5), 0)
    #  seq.sort(key=operator.attrgetter("key"), reverse=True)
    #  seq.sort(key=lambda a: (a.key, d[a]), reverse=True)
    return list2str(rest)

############################################################
# random testcase
def rand_testcase():
    raise NotImplementedError, "Please implement rand_testcase() by yourself."
    #N = randint(5,1000)
    #N = randrange(5,1000+1,1)
    #chars = [chr(x) for x in range(ord('a'), ord('z')+1)] + [' ']
    #Ls = [choice(chars) for x in range(1000)]
    #Ds = sample(range(10), 2)
    #del chars, x
    return locals() 

############################################################
#
#  Note:
#    Following lines are come from the template
#    which is distributed on the blog article:
#      http://d.hatena.ne.jp/over80/20111007/1318002963
#
#    There are no problem specific codes from here to the end of file.
#

opt_single = True
opt_random = 0
opt_call_all_solvers = False
opt_quiet = False

try:
    opts, args = getopt.gnu_getopt(sys.argv[1:], "mr:cq")
    for o, a in opts:
        if o == '-m': opt_single = False
        elif o == '-r': opt_random = int(a)
        elif o == '-c': opt_call_all_solvers = True
        elif o == '-q': opt_quiet = True
        else: raise Exception("unknown option(%s)" % (o,))
    outfilename = None
    if not opt_random:
        infilename = args[0]
        outfilename = os.path.splitext(infilename)[0] + ".out"
        infile = open(infilename, 'r')
except Exception, e:
    print >>sys.stderr, e
    print >>sys.stderr, 'Usage: %s [-c] [-q] [-m] <inputfile>' % (sys.argv[0])
    print >>sys.stderr, '       %s [-c] [-q] [-m] -r num' % (sys.argv[0])
    exit(1);

if opt_random:
    precond = dict()
    num_tests = opt_random
    get_testcase = rand_testcase
else:
    precond = read_precond()
    num_tests = precond.pop("NoTC")
    get_testcase = read_testcase

# set output target
outfile = [sys.stdout,]
if outfilename != None and not opt_call_all_solvers:
    outfile.append(open(outfilename, 'w'))
if opt_quiet:
    sys.stdout = open('/dev/null', 'w')

# prepare testcase iterator
def read_testcases(num_tests):
    for i in range(num_tests):
        yield get_testcase()
testcases = read_testcases(num_tests)

# prepare solve iterator
def call_solve(q):
    d = dict(precond)
    d.update(q)
    if opt_call_all_solvers:
        res = set()
        for name, func in solvers:
            r = func(**deepcopy(d))
            print >>sys.stderr, "SOLVER:", name, "  \tANSWER:", r
            res.add(r)
        if len(res) != 1:
            raise RuntimeError, "ansers does not match"
        return r
    else:
        return solve(**d)
def call_solve_trap(q):
    try:
        return call_solve(q)
    except KeyboardInterrupt:
        pass
if opt_single:
    res = imap(call_solve, testcases)
else:
    p = multiprocessing.Pool()
    res = p.imap(call_solve, testcases)

# solve & print
for (i, r) in enumerate(res, start=1):
    for f in outfile:
        print >>f, "Case #%d: %s" % (i, str(r))