AtCoder Beginner Contest 314 C - Rotate Colored Subsequence

AtCoderをやり始めた。 全然問題が解けないので、復習する。

問題

atcoder.jp

問題文が長いので、理解するまでに結構な時間を使った。

私的考察

1. 入力した文字列 $S$ に対して、$i$ごとに文字を抽出する。
2. 抽出した$i$ごとに中の文字をシフトさせる。
3. 再度$i$ごとに値を取り出し、文字列を再構築。

みたいにすればOKだと思い込み実装。

結果： めっちゃメモリ使ってた。。。。orz

解き方

これは、$i$ごとに文字を抽出するのではなく、添字を抽出するのがポイント。

加えて、配列ではなくmapを使うことで、for文が3個くらい減らせる。

package main

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

func main() {
    in := bufio.NewReader(os.Stdin)
    out := bufio.NewWriter(os.Stdout)
    defer out.Flush()

    var N, M int
    fmt.Fscan(in, &N, &M)

    var S string
    fmt.Fscan(in, &S)
    ss := strings.Split(S, "")

    mp := map[int][]int{}
    for i := 0; i < N; i++ {
        var c int
        fmt.Fscan(in, &c)
        mp[c] = append(mp[c], i)
    }

    for i := 1; i <= M; i++ {
        if len(mp[i]) <= 1 {
            continue
        }
        last := ss[mp[i][len(mp[i])-1]]
        for j := len(mp[i]) - 1; j > 0; j-- {
            ss[mp[i][j]] = ss[mp[i][j-1]]
        }
        ss[mp[i][0]] = last
    }
    fmt.Fprintln(out, strings.Join(ss, ""))
}

あと、今回はいつも使用しているbufio.Scannerを使うと、うまく読み取れない問題があ出てきたので、 bufio.NewReaderとbufio.NewWriterを使用した。

うーん、このあたりも理解しておかないとこの先辛いな。。。