COLOCON 2018-Final: A - ファイティング・タカハシ
問題
解説
'B'で区切られた'A'について、それぞれの'A'の連結成分においてスコアを算出する。
ここでの値が大きいので、当然を個連結して解くことはできない。
の末尾との先頭が連結するかどうかとに着目する必要がある。 またすべてが'A'であるとき、個の'A'が存在し、スコアの算出方法が異なるため、別途取り扱う必要がある。
つまり、以下の3つに場合分けできれば良い。
- のすべての要素が'A'
- の末尾は'A' かつ の先頭が'A'
- の末尾は'B' または の先頭が'B'
'B'は区切り文字としてしか機能しないため、を'A'で作られる文字列の配列に変換し、, , に分けてスコアの算出を考えた。
auto make_A_vec(string const& S) { vector<string> ret; string curr; for (auto c: S) { if (c == 'B') { if (curr.size()) { ret.push_back(curr); curr.clear(); } } else { curr += 'A'; } } if (curr.size()) { ret.push_back(curr); } return ret; } template<typename InputIterator> auto sum_A(InputIterator first, InputIterator last) { int64_t sum = 0; while (first != last) { sum += first->size() * (first->size() + 1) / 2; first++; } return sum; } int main() { int N; cin >> N; string S; cin >> S; { auto t = S; t.erase(std::unique(t.begin(), t.end()), t.end()); if (t.size() == 1) { if (t[0] == 'B') { cout << 0 << "\n"; return 0; } int64_t x = S.size() * N; cout << x * (x + 1) / 2 << "\n"; return 0; } } auto v = make_A_vec(S); if (S.back() == 'A' && S[0] == 'A') { auto X = (int64_t)v[0].size() * (v[0].size() + 1) / 2; auto sumY = v.size() > 2 ? sum_A(v.begin() + 1, v.end() - 1) : 0; auto Z = (int64_t)v.back().size() * (v.back().size() + 1) / 2; auto ZX = (int64_t)(v.back().size() + v[0].size()) * ((v.back().size() + v[0].size()) + 1) / 2; cout << (X + sumY) + (N - 1) * (ZX + sumY) + Z << "\n"; } else { int64_t sum = 0; for (auto e: v) { sum += (int64_t)e.size() * (e.size() + 1) / 2; } cout << sum * N << "\n"; } } /* // 実装力のなさから単体テストを書く羽目になった vector<vector<string>> vs = { {"A", "A", "A"}, {"A", "A", "A", "A"}, {"A", "A", "AA", "A"}, {"A", "A", "AAA", "A"}, {"A", "AA", "AAA", "AAAA", "A"}, }; vector<int64_t> res_vs = { 1, 2, 4, 7, 3 + 6 + 10 }; rep(i, res_vs.size()) { assert(sum_A(vs[i].begin() + 1, vs[i].end() - 1) == res_vs[i]); } return 0; */ /* assert((make_A_vec("A") == vector<string>{"A"})); assert((make_A_vec("AA") == vector<string>{"AA"})); assert((make_A_vec("ABA") == vector<string>{"A", "A"})); assert((make_A_vec("ABABB") == vector<string>{"A", "A"})); assert((make_A_vec("ABABBAA") == vector<string>{"A", "A", "AA"})); return 0; */