AtCoder diverta 2019 Programming Contest C - AB Substrings

AtCoderに取り組んだ記録です。
AtCoder ProblemsのRecommendationsからピックアップして解いています。

問題ページはこちらです。
Point: 400
Difficulty: 848

考えたこと

N <= 10の4乗なので、すべての順序を確認するのは間に合わないと思います。

連結することでつくれる"AB"の個数
A->Bと連続する場合のみが関心の対象であるため、連結について考慮すべき文字列は以下の3つのパターンに分かれます。
1. 末尾が"A"の文字列（"ABCA"など）
2. 先頭が"B"の文字列（"BAD"など）
3. 先頭が"B"で末尾が"A"の文字列（"BMHQUVA"など）

3の文字列が存在する場合、可能な限り3の文字列を連続させることで、連結による"AB"の個数を増やすことができます。連結でれきるABの個数は "3の文字列の個数 - 1" 個になります。

3の文字列群で作成した文字列の前に1の文字列をつなぐことで、"AB"を1個増やせます。
同様に、後ろに2の文字列をつなぐことで、1個増やせます。

残った1と2で"AB"をつくるため、1と2のうち小さい方の個数分、"AB"を増やせます。

各文字列中に含まれる"AB"の個数
こちらは文字列中に含まれる"AB"の個数を数えればOKです。

上記の2つの方法で数えた個数を合算すれば解答になります。

なんだかまわりくどい方法な気がしますが、これをそのまま実装することでACできました。

提出したコード

提出ページ

package main

import (
	"bufio"
	"fmt"
	"os"
	"strings"
)

func main() {
	r := bufio.NewReader(os.Stdin)
	var n int
	fmt.Fscan(r, &n)
	ss := make([]string, n)
	for i := 0; i < n; i++ {
		fmt.Fscan(r, &ss[i])
	}

	suffixACount := 0
	prefixBCount := 0
	aToBCount := 0
	abCount := 0
	for _, s := range ss {
		abCount += strings.Count(s, "AB")

		if s[0] == 'B' && s[len(s)-1] == 'A' {
			aToBCount++
		} else if s[0] == 'B' {
			prefixBCount++
		} else if s[len(s)-1] == 'A' {
			suffixACount++
		}
	}

	if aToBCount > 0 {
		abCount += aToBCount - 1
		if suffixACount > 0 && prefixBCount > 0 {
			abCount += 2
			abCount += min(suffixACount-1, prefixBCount-1)
		} else if suffixACount > 0 {
			abCount++
		} else if prefixBCount > 0 {
			abCount++
		}
	} else {
		abCount += min(suffixACount, prefixBCount)
	}

	fmt.Println(abCount)
}

// Utility

func min(x, y int) int {
	if x < y {
		return x
	}
	return y
}

func max(x, y int) int {
	if x > y {
		return x
	}
	return y
}

感想

・3の文字列の取り扱いに気づけてよかった