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

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

AtCoder ABC 137 C - Green Bin (Go)

irisruneです。前回のコンテストは色々あり2完で-91(1292→1201)と大変冷えました。

問題

atcoder.jp

解き方自体は発想さえできればといったところですが、言語のソート仕様によっては実装が面倒になります。

実はこの提出コードはAC時のものではありません。どこに問題があるか探してみるのもいいですね。

package main

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

var sc *bufio.Scanner

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

type ByRune []rune

func (r ByRune) Len() int {
    return len(r)
}

func (r ByRune) Swap(i, j int) {
    r[i], r[j] = r[j], r[i]
}

func (r ByRune) Less(i, j int) bool {
    return r[i] < r[j]
}

func main() {
    sc = bufio.NewScanner(os.Stdin)
    sc.Split(bufio.ScanWords)
    n := nextInt()

    strMap := make(map[string]int)
    for i := 0; i < n; i++ {
        var s string
        fmt.Scan(&s)
        rs := []rune(s)
        sort.Sort(ByRune(rs))
        sSorted := string(rs)
        _, ok := strMap[sSorted]
        if ok {
            strMap[sSorted]++
        } else {
            strMap[sSorted] = 1
        }
    }

    ans := 0
    for _, v := range strMap {
        ans += v * (v - 1) / 2
    }
    fmt.Println(ans)
}

2つの文字列の組がアナグラムの関係になっているかを判断するには文字種ごとの文字数を考えるのが直感的です。しかし、全ての文字列のペアに対してアナグラムの関係になっているかチェックすると計算量がO(N^2)になってしまいます。

そこで辞書型(あるいは連想配列)を用いることにより、アナグラムの関係になっている文字列のグループを作ることを考えます。最終的な答えは、各グループの要素数をM_jとおくと\sum M_j(M_j - 1)/2と求められます。

文字種ごとの文字数を考える場合でも、26桁の11進数に変換するなどの工夫をすれば上記の方法に落とし込むことができます。ただ、各文字列を文字単位でソートする方がスマートに見えると思います。なお、計算量は一般的にはソートする方が大きいですが、文字列長が10で固定のためソートする方が小さくなります。

ここまで来ればあとは実装するだけなのですが、Goのソート仕様が大きく足を引っ張ってしまいます。Goで構造体のリストをソートする際に散々苦しめられてきているのはご存知かもしれませんが、実は組み込み型についてもInt,Float64,String以外の型についてはソート関数が標準で実装されていません(正確には、ソートインタフェースへの変換関数が実装されていません)。

そのため、文字(Rune型)の配列を新たに型として定義し、定義した型についてソートインタフェースの実装を行ってからソートを行う必要があります。そもそもRune型の存在を知っていないとどうにもならないので文字列仕様とソート仕様のダブルパンチを受けている格好です。

ところで上記コードはAC時のものではないと述べましたが、なんと全てのケースについてWAが出てしまいます。D問題でTLEして後に引けなくなったところにC問題の全WAで心が折れてしまったのが大寒波の正体でした。

原因としてはsc.Scan()を一度実行してしまうと以降のfmt.Scan()で取得される文字列が全て空になってしまうというものです。そのため、nの取得もfmt.Scan()で行う(あるいは、sの取得をsc.Scan()→sc.Text()で行う)ようにすればACとなります。

バージョンの差異による問題か、手元で実行した時には出力結果が正常だったので焦りもあってコンテスト中に原因が突き止められませんでしたが、後でAtCoderのコードテストで試すと出力が0であったために気付くことができました。

雑記

  • 今回のようなケースを考えるとやはり手元の環境とAtCoder上の環境とで言語バージョンは揃えた方がいいのですが、Go周りの(公式)ツールが正常にインストールされないという問題があるのが厳しいです。
    • 言語アップデートでAtCoder側がバージョンアップしてくれればあまり考えることもないのですが、なかなかアップデートが実施されないですね。