前回はこちらから。

最近、問題を解いていて、基本的なアルゴリズムの使いどころや実装を知らないことに気づいた。このままただ問題解き続けても、あまり成長しそうにないなぁと。
なので、そういった基本的なところの勉強ができるかなと思って蟻本を買ってみることに。

まだ序盤の全探索ぐらいしか読んでいないけども、もっと早く買っておけばよかったなぁと。頻出の基本的アルゴリズムやどの問題に対して、アルゴリズムを用いるべきかなどがわかるから、結構便利だし考え方のコツがつかめる。

これ読んでさらにアルゴリズムに強くならないと。

せっかく蟻本買ったので、記載されているアルゴリズムを用いて解く問題を重点的に解きたいと思い、どっかにまとめられてないかなと探してみた。
そうしたら以下のサイトがヒット！。蟻本に記載されているアルゴリズムに対応したAtCoderの問題をまとめてもらえていたので、それを解いていくことに。

ここから、ABC問題を選んで解いていこうかなと思う。

というわけで、今回は全探索を用いるABC045Cを解いていこう。

ABC045Cの問題

問題文
1 以上 9 以下の数字のみからなる文字列 S が与えられます。
この文字列の中で、あなたはこれら文字と文字の間のうち、いくつかの場所に + を入れることができます。
一つも入れなくてもかまいません。
ただし、+ が連続してはいけません。このようにして出来る全ての文字列を数式とみなし、和を計算することができます。
ありうる全ての数式の値を計算し、その合計を出力してください。

制約
- 1≤|S|≤10
- S に含まれる文字は全て 1 〜 9 の数字

アルゴリズム

「ありうるすべての値を計算し～」って書いてあるし、全探索を行えって言われているな。各数字の間に＋を入れるか入れないかのパターンすべてを求めて計算すればよい感じかな。

というわけで、深さ優先探索を用いて全パターンの計算結果を取得することに。アルゴリズムは以下のような感じに。

1. ルートノードを生成し、深さ優先探索を開始する。
2. ノードは、探索の深さに対応した各数字の間に＋を入れる入れないの2つのパターンの子ノードを生成する。
3. 一番下まで探索したら、親ノードに現在の式(文字列に＋を入れたもの)の計算結果を返す。
4. 生成した2つの子ノードの計算結果を足して、親ノードに返す。
5. ルートノードに値が返ってきたら、それを答えとする。

深さ優先探索のノードの処理の擬似コードはこんな感じ。

// s: 与えられた文字列, flags: +を入れるかのフラグ, depth: 残りの探索の深さ
func dfs(s, flags, depth) int {
    // 末端ノードまで来たら、計算結果を返す
    if depth == 0 {
	return calc(s, flags)
    }

    res := dfs(s, flags, depth-1) // +を入れない場合の子ノード
    res += dfs(s, flagOn(flags, depth), depth-1) // ＋を入れる場合の子ノード
    return res
}

実際には、＋を入れるかどうかのフラグはbit列を用いて管理することに。bitが1ならば、その位置には＋を入れて、0なら入れないみたいに。

例）
　文字列：12345
　フラグ：0101
結果）
　式：12 + 34 + 5

上記の例では、2と3、4と5の間に対応するbitは1となっているので、その間に＋を入れている。

bit列で管理するため、擬似コード上での flagOn() は以下のようになった。

// flagOn()の処理を模倣するbit演算 (flags -> f, depth -> d)
// 左からdepth+1とdepth番目の数字の間に＋を入れる
f | (1 << (d-1)) 

正直言って可読性が落ちてしまうが、今回は処理のシンプルさを優先することに。

あとは、得た式から結果を計算する calc() だけど、これは上記の例を再現すればよいので、以下のような感じに。

func calc(s string, f int) int {
    v, j := 0, 0
    for i := range s {
	if f&1 == 1 { // フラグが立っているか
            n, _ := strconv.Atoi(s[j : i+1]) // +のところまで文字列を切り取り数値化
	    v += n // 数値化した値を足す
	    j = i + 1
	}
	f >>= 1 // 次のフラグへシフト
    }
    n, _ := strconv.Atoi(s[j:])
    v += n
    return v
}

さぁこれで必要な処理はできたから、入出力つけて完成だ～！

解答

最終的には以下のようなコードに。

package main

import (
    "fmt"
    "strconv"
)

func main() {
    var s string
    fmt.Scan(&s)

    fmt.Println(dfs(s, 0, uint(len(s))-1))
}

func dfs(s string, f int, d uint) int {
    if d == 0 {
	return calc(s, f)
    }
    return dfs(s, f, d-1) + dfs(s, f|(1<<(d-1)), d-1)
}

func calc(s string, f int) int {
    v, j := 0, 0
    for i := range s {
	if f&1 == 1 {
            n, _ := strconv.Atoi(s[j : i+1])
	    v += n
	    j = i + 1
	}
	f >>= 1
    }
    n, _ := strconv.Atoi(s[j:])
    v += n
    return v
}

今回の問題は、典型的な全探索問題で、本で読んだことがそのまま使えたりしてスムーズ＆楽しく解くことができた。

次回以降も、蟻本読みながら問題解いて、いろいろな解き方に慣れていくぞー