けんちょんの競プロ精進記録

競プロの精進記録や小ネタを書いていきます

AtCoder AGC 040 A - >< (茶色, 300 点)

AtCoder AtCoder300点 AGC-A 茶色diff 数列 >< 左右からそれぞれ走査する Greedy 最適化テク:自明な上界が最適解 ジグザグ

いかに上手に実装するか！

問題概要

長さ N−1 の ">" と "<" のみかななる文字列 S が与えられる。

長さ N の非負整数列 a_{1}, a_{2}, \dots, a_{N} は， すべての i について次の条件をみたすとき、良い非負整数列と呼ぶ。

  • S_{i} = "<" のとき： a_{i} \lt a_{i+1}
  • S_{i} = ">" のとき： a_{i} \gt a_{i+1}

良い非負整数列の要素の総和としてありうる最小の値を求めよ。

制約

  • 2 \le N \le 5 \times 10^{5}

考えたこと

とりあえず、

"><"

となっている箇所については、数値は「0」としてよい。">>><<" という感じだったら、

(3, 2, 1, 0, 1, 2)

とうふうになる。問題になるのは、"<>" となっている箇所。この部分については、

  • 左から登ってきた階段 (0, 1, 2, ..., h1)
  • 右から登ってきた階段 (h2, ..., 2, 1, 0)

のうちの高い方を採用すればよい。

実装

一番楽な実装は、左右それぞれから1周ずつ以下のことをする感じだと思う。

  • 左から見ていって「<」だったら、chmax(現在の値, 直前の値 + 1) とする
  • 右から見ていって「>」だったら、chmax(現在の値, 直後の値 + 1) とする

計算量は O(N) となる。

コード

#include <bits/stdc++.h>
using namespace std;
template<class T> inline bool chmax(T& a, T b) { if (a < b) { a = b; return 1; } return 0; }

int main() {
    string S;
    cin >> S;
    int N = S.size() + 1;
    vector<int> A(N, 0);
    for (int i = 0; i < N-1; ++i) {
        if (S[i] == '<') chmax(A[i+1], A[i]+1);
    }
    for (int i = N-2; i >= 0; --i) {
        if (S[i] == '>') chmax(A[i], A[i+1]+1);
    }
    long long res = 0;
    for (auto v : A) res += v;
    cout << res << endl;
}