あるエンジニアのAtCoder奮闘記

東京都港区にあるアミフィアブル株式会社のエンジニアが、AtCoderで解いた問題について振り返ったりしていく会社公認のブログです。

AtCoder AGC 040 A - >< (Go)

irisruneです。今週のコンテストは1完でまた水パフォでした。

問題

atcoder.jp

見かけがかなりわかりにくく、アルゴリズムが組みにくい問題だと思います。

package main

import (
    "bufio"
    "fmt"
    "os"
    "strconv"
)

var sc *bufio.Scanner

func nextStr() string {
    sc.Scan()
    return sc.Text()
}

func nextInt() int {
    sc.Scan()
    i, e := strconv.Atoi(sc.Text())
    if e != nil {
        panic(e)
    }
    return i
}

func main() {
    sc = bufio.NewScanner(os.Stdin)
    sc.Buffer(make([]byte, 0), 1000000001*3)
    sc.Split(bufio.ScanWords)
    s := nextStr()
    ans := 0
    mode := '>'
    span, peak := 0, 0
    for _, r := range s {
        switch {
        case r == mode:
            span++
        case r == '>':
            peak = span
            ans += (span * (span + 1)) / 2
            span = 1
            mode = r
        case peak >= span:
            ans += (span * (span - 1)) / 2
            span = 1
            mode = r
        case peak < span:
            ans += (span - peak)
            ans += (span * (span - 1)) / 2
            span = 1
            mode = r
        }
    }
    switch {
    case mode == '<':
        ans += (span * (span + 1)) / 2
    case peak >= span:
        ans += (span * (span - 1)) / 2
    case peak < span:
        ans += (span - peak)
        ans += (span * (span - 1)) / 2
    }
    fmt.Println(ans)
}

簡単なケースについて各A_iの最小値を考えてみます。例えば<<<>>という入力に対しては、0,1,2,3,1,0となり、一方で<<>>>という入力に対しては、0,1,3,2,1,0となります。ここから考えると、各要素の最小値について以下のように示せます。

  1. 『要素の左側が>あるいは数列の左端』かつ『要素の右側が<あるいは数列の右端』である場合、最小値は0となる。
  2. 『要素の左側が>あるいは数列の左端』かつ『要素の右側が>』である場合、最小値は右隣の要素の最小値より1大きな値となる。
  3. 『要素の左側が<』かつ『要素の右側が<あるいは数列の右端』である場合、最小値は左隣の要素の最小値より1大きな値となる。
  4. 『要素の左側が<』かつ『要素の右側が>』である場合、最小値は左隣の要素の最小値と右隣の要素の最小値とで大きい方の値より1大きな値となる。

このうち1.から3.までは簡単に求められ、4.も値を一時的に記憶しておくなどすれば求めることができます。計算量はO(N)です。

雑記

  • 公式解説にあるように、『要素の左側に連続する<の数』と『要素の右側に連続する>の数」とで大きい方の値を取ると、上記1.から4.までの性質をすべて包含するためより簡潔に求めることができます。
    • ここに気づくことができればABC-C問題相当と言えるかもしれませんが、一般的に見てABC-D問題相当くらいの難易度になると思います(6問形式のABC基準です)。