AOJでHaskell: 1188 階層民主主義 (Hierarchical Democracy)

こんにちはtatsyです。

今日はHaskellの勉強を兼ねてAOJの関数型向きっぽい問題をやってみます。やる問題は1188番の階層民主主義という問題です。

問題の内容


入力は[[123][4567][89]]のようなかっこで区切られた文字列です。この中身は有権者の人数を表していて、この半分以上を取ればこの地区で勝利できます。

上の入力例ではかっこで囲まれた数字が3個ありますが、候補者はこの3個のうちの半数以上で勝利すれば当選です。なので、上の例では123の半分である62と89の半数である45の合計である107が当選に必要な最小票数ということになります。

上の例では階層が2つだけでしたが、実際の入力はこれが複数の層になっています。

解法


問題自体はそれほど難しくないので解法はすぐ思いつきそうです。単純に考えれば、かっこで囲まれた範囲を再帰的に展開して、下から計算をしていけばいいだけです。

なので、C++とかでやるのであれば、かっこで囲まれた同階層の内容をsplitする関数を書いて、splitされた各要素を計算する関数を呼び出せばよさそうです。

それで書いた解答がこちらになります。


#include <iostream>
#include <string>
#include <vector>
#include <algorithm>
using namespace std;

// 文字列を分割する
vector<string> split(string str, string pat) {
    int pos;
    vector<string> ret;
    while ((pos = str.find(pat)) != string::npos) {
        ret.push_back(str.substr(0, pos));
        str = str.substr(pos + pat.size());
    }
    ret.push_back(str);
    return ret;
}

// 文字列の前後の文字を1文字ずつ削る
string strip(string str) {
    return str.substr(1, str.size() - 2);
}

// 文字列で表された数字をパースする
int parse(string str) {
    int ans = 0;
    for (int i = 0; i < str.size(); i++) {
        ans = ans * 10 + (str[i] - '0');
    }
    return ans;
}

int solve(string str) {
    if ('0' <= str[0] && str[0] <= '9') {
        // 入力が数字ならパースする
        return (parse(str) + 1) / 2;
    } else {
        // そうでなければ文字列を配列の区切れで分割する
        vector<string> ss;
        while (!str.empty()) {
            int acc = 0;
            for (int i = 0; i < str.size(); i++) {
                if (str[i] == '[') {
                    acc += 1;
                } else if (str[i] == ']') {
                    acc -= 1;
                }

                if (acc == 0) {
                    ss.push_back(str.substr(0, i + 1));
                    str = str.substr(i + 1);
                    break;
                }
            }
        }

        // 各配列に対して再帰的に関数を呼び出す
        vector<int> ns;
        for (int i = 0; i < ss.size(); i++) {
            ns.push_back(solve(strip(ss[i])));
        }

        // 必要な票数は下から半分の数字の合計
        sort(ns.begin(), ns.end());
        int ans = 0;
        int m = (ns.size() + 1) / 2;
        for (int i = 0; i < m; i++) {
            ans += ns[i];
        }
        return ans;
    }
}

int main(int argc, char** argv) {
    int n;
    cin >> n;
    for (int i = 0; i < n; i++) {
        string str;
        cin >> str;
        printf("%d\n", solve(strip(str)));
    }
}

単純ですね。でも、これを関数型でやろうとすると、案外に文字列を分割する処理が面倒そうです。

上のC++のコードでは、かっこを左から順にみていって[が出てきたら1を足して、]が出てきたら1を引くという処理を行い、この数が0になったところで文字列を分割しています。

HaskellでもSTモナドとか使えばできると思いますが、あまりエレガントじゃないですし、そもそも関数型の勉強にならないので、再帰関数で考えてみます。

基本的な考え方として[が出てきたら、新しい階層に入ったと見なして、再帰関数を1つ下に下ります。各再帰関数は現れた数字を保存しておくstackのようなものを持たせておいて、]が現れたらstack内に積まれた数字の小さいもの半分の合計を上に返すようにします。

と、さらっと書きましたが、結構考えるのに時間かかりました(笑)。実装したコードは次のようになります。

import Control.Monad
import Control.Applicative
import Data.List
import Debug.Trace

-- 先頭から始まる数字をパースして、数字と残りの文字列を返す
parse :: String -> (Int, String)
parse xs =
    let (prev,t:rst) = break (== ']') xs
        x = read prev
    in ((x + 1) `div` 2, rst)

-- 得られている数字の配列に対して、小さい方から半分の合計を取る
calc :: [Int] -> Int
calc xs =
    let m = ((length xs) + 1) `div` 2
    in sum $ take m $ sort xs

-- 問題を解く再帰関数
solve :: [Int] -> String -> (Int, String)
solve stk "" = (calc stk, "")
solve stk (c:cs)
    -- 開きかっこなら一段下に下がる
    | c == '['  =
        let (x, rst) = solve [] cs
        in solve (x:stk) rst
    -- 閉じかっこなら今スタックに積んである数を計算する
    | c == ']'  =
        let x = calc stk
        in (x, cs)
    -- 数字なら普通にパースする
    | otherwise = parse (c:cs)

main :: IO ()
main = do
    n <- readLn
    ls <- replicateM n getLine :: IO [String]
    mapM_ print $ map (fst . solve []) ls

再帰関数を工夫した分だけコードの量は短くなってますね。

自分なりには結構エレガントなんじゃないかと思います。が、C++のやり方の方が直感的に分かりやすいというのは否定できないですね。

こういう再帰関数を使った書き方がささっと思いつくようになるといいなぁと思います。

というわけで内容がないですが、これで今回の記事を終わります。

お読みいただきありがとうございました。