ローラーコースターでいくら稼げるか

タイトルで釣られた人はごめんなさい。プログラミングのアルゴリズムの話題です。

最後の30分ぐらいしか出れなかったけれど、http://community.appliplanet.co.jp/groups/google-code-jam-2011-%E5%8B%89%E5%BC%B7%E4%BC%9A/blog/2011/04/06/%E3%82%B0%E3%83%BC%E3%82%B0%E3%83%AB%E3%82%B3%E3%83%BC%E3%83%89%E3%82%B8%E3%83%A3%E3%83%A0%E5%8B%89%E5%BC%B7%E4%BC%9A-%EF%BC%94-%EF%BC%95-%E7%81%AB-%E3%81%AE%E5%86%85%E5%AE%B9 に参加してきました。その場でそそくさと実装してできたかな？と思ってたら、今日になってサンプル入力(small)の出力を提出してみる(答え合わせ)とバグってて、難しいなーと実感した次第です。

問題文は以下になります。

GCJ2010 QR C. Theme Park テーマパーク

愚直に実装してみたのが以下です。入力は指定フォーマットのテキストファイルなので、ファイル読み込みの実装は適当ですが、主題とは関係ないのであまり気にしないでください(^ ^;;

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

import sys

def solve_gcj2010_qr_c_theme_park(testno, test_case):
    R, K, N = map(int, test_case[0].split())
    g = map(int, test_case[1].split())
    sales = 0
    for _ in range(R):
        roller_coaster = []
        current_num = 0
        over_flag = False
        for num, i in enumerate(g):
            if current_num + i <= K:
                current_num += i
                roller_coaster.append(i)
            else:
                # start coaster
                over_flag = True
                break
        if over_flag:
            g = g[num:] + roller_coaster
        sales += current_num
    print "Case #{0}: {1}".format(testno, sales)

def read_input(file_name):
    t = []
    with open(file_name) as f:
        for lineno, line in enumerate(f):
            if lineno != 0:
                t.append(line)
                if lineno % 2 == 0:
                    yield t
                    t = []

def main():
    for num, test_case in enumerate(read_input(sys.argv[1])):
        solve_gcj2010_qr_c_theme_park(num + 1, test_case)

if __name__ == "__main__":
    main()

当初は全てのグループがローラーコースターに乗れるときの処理を考慮してなくてバグってました。over_flag というフラグとか設けてごにょごにょしてます。普段は、こういうループ内でフラグ使うような実装はしないのですが、お仕事じゃないし、途中で面倒になってもういいやと実装してしまいました(^ ^;;

この愚直な実装だと、ローラーコースターの回転数 R 、グループ数 N 、テストケース T が与えられると、

R x N x T

の計算量が必要で、問題文にある巨大な値がセットされているサンプル入力(large)だと、計算にもの凄く時間がかかってしまいます。そこでアルゴリズムをちゃんと考えましょうねというお話になります。勉強会の最後に解法の解説もしていました。それによると、この R を減らすか、N を減らすかの2つの解法があるようです。先のページに解答例もあります。

この解答例の中の

グループ数の最大値は小さいので、各グループ i に対して、列の先頭がグループ i だった時、その回に何人が乗れるか、次の回の先頭グループはどこか, を予め計算しておく。
GCJ 2010 Qualification Round C.Theme Park - hotoku とは

が、個人的にとても分かり易かったので、この方針を参考に再実装してみました。

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

import sys

def calculate_num_from_group_position(g, N, K):
    results = []
    for pos in range(N):
        _g = g[pos:] + g[:pos]
        current_num = 0
        for _pos, i in enumerate(_g):
            if current_num + i <= K:
                current_num += i
            else:
                break
        next_pos = (pos + _pos) % N
        results.append((current_num, next_pos))
    return results

def solve_gcj2010_qr_c_theme_park(testno, test_case):
    R, K, N = map(int, test_case[0].split())
    g = map(int, test_case[1].split())
    pos_to_num = calculate_num_from_group_position(g, N, K)
    sales = 0
    next_pos = 0
    for _ in range(R):
        current_num, next_pos = pos_to_num[next_pos]
        sales += current_num
    print "Case #{0}: {1}".format(testno, sales)

def read_input(file_name):
    t = []
    with open(file_name) as f:
        for lineno, line in enumerate(f):
            if lineno != 0:
                t.append(line)
                if lineno % 2 == 0:
                    yield t
                    t = []

def main():
    for num, test_case in enumerate(read_input(sys.argv[1])):
        solve_gcj2010_qr_c_theme_park(num + 1, test_case)

if __name__ == "__main__":
    main()

私の環境 (Core i7 2GHz/8GB) では、サンプル入力(large) が数分で処理できました。Python の for 文は for (i=3; i

2011/4/8 追記コメントから for ループを while ループにしてみました

def calculate_num_from_group_position(g, N, K):
    results = []
    _g = g * 2
    for pos in range(N):
        current_num = 0
        _g_max = pos + N
        while pos < _g_max:
            if current_num + _g[pos] <= K:
                current_num += _g[pos]
            else:
                break
            pos += 1
        next_pos = pos % N
        results.append((current_num, next_pos))
    return results

ループで使用している pos を再利用しているのが気持ち悪いですが、これでも同じですね。数値のコピーをどうやるのか分からなかった(> <)

2011/4/8 追記 コメントから for ループを while ループにしてみました

2011/4/8 追記コメントから for ループを while ループにしてみました