気分転換に競技プログラミング: Google Code Jam Qualification Round 2012: Problem B. Dancing With the Googlers 改

前回の続き。
前回はとりあえず全探索的な感じで書いたので、アルゴリズム効率が悪くLargeなケースだと解けずにいたが、先日一緒にこの問題について悩んでいた後輩がその翌日会った矢先に「解けました！」と嬉しそうにPerlのコードを見せてきた。そのコードをPythonで書き直したものを下に晒しておく。(きっと需要はないだろうが)

solve.py

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

"""
Google Code Jam
Qualification Round 2012
Problem B. Dancing With the Googlers
http://code.google.com/codejam/contest/1460488/dashboard#s=p1
"""

import sys

def read():
    f = sys.stdin
    num = int(f.next().strip())
    return [f.next().strip() for i in range(num)]

def solve(case):
    line = case.split(" ")
    N = int(line.pop(0))
    S = int(line.pop(0))
    P = int(line.pop(0))
    T = map(int, line)
    count = other = 0
    if P == 0:
        return len(T)
    elif P == 1:
        for t in T:
            if t >0: count += 1
        return count
    else:
        for t in T:
            if t >= 3*P-2:
                count += 1
            elif (t == 3*P-3 or t == 3*P-4):
                other += 1
        if other >= S:
            return S + count
        else:
            return other + count

def main():
    for i, case in enumerate(read()):
        answer = solve(case)
        print "Case #%d: %d" % (i+1, answer)

if __name__ == "__main__":
    main()

それにしても、その後輩が直接何かに熱中した姿を見たことがなかったのもあり、競技プログラミングを通してそんな無邪気な一面を見せてくれてこちらも嬉しかった。修了あとわずかだが(無事に修了できる気はしないが)、研究室で定期的に競技プログラミングの問題に取り組めるといいなと思っている今日この頃である。そして今日も修論から現実逃避をしたのであった。