Go言語で挑む競プロ #13

前回はこちらから。

昨日、久しぶりに予定を合わせることができて、Atcoderのコンテスト（ABC137）に参加することができた！
久々のコンテストだったけども時間内に何とかA、B、C問題まで解くことができた。一度のミスなく提出できて結構満足。D問題にも取り掛かりたかったけども、それは今度のお楽しみに取っておこう。

さぁというわけで、本日2つ目のnoteだけど、解いたABC137のC問題を書いていこうかなと思う。

ABC137Cの問題

問題文
文字列 a に含まれる文字を何らかの順序で並べることで得られる文字列を a の アナグラム と呼びます。
例えば、greenbin は beginner のアナグラムです。
このように、同じ文字が複数回現れるときはその文字をちょうどその回数だけ使わなければなりません。
N 個の文字列 s1,s2,…,sN が与えられます。
それぞれの文字列は長さが 10 で英小文字からなり、またこれらの文字列はすべて異なります。
二つの整数 i,j (1≤i<j≤N) の組であって、si が sj のアナグラムであるようなものの個数を求めてください。

制約
- 2≤N≤10^5
- si は長さ 10 の文字列である。
- si の各文字は英小文字である。
- s1,s2,…,sN はすべて異なる。

アルゴリズム

アナグラムかどうかは、ソートすればわかるしあとは数え上げるだけで解けそうかな。なので、以下のようなアルゴリズムを考えてみた。

1. 各文字列をソートしてアナグラムかどうかをはっきりさせる
2. ソート後に同じ文字列となった文字列（アナグラム）の数をカウントする
3. カウントした結果から、各文字列のアナグラムの組の数を求める
4. すべての文字列のアナグラム数を合計して解答とする

とてもシンプルなアルゴリズムだね。この手順の中で少し不明確なのは手順3のアナグラムの組の数を求めるところかな。

今回の問題は、i < j で　si が sj のアナグラムである場合の数を解答するので、j < i のときは数えてはいけない。そのため、アナグラムの組の数は、手順2で求めているアナグラム数ｎを用いて、n-1C2（組み合わせ）で求めることができる。

組み合わせ計算を実装するのはちょっとめんどくさいけど、実はnC2 は n + (n -1) + (n-2) + ... + 1 で求めることができるので意外と簡単に求めることができる。以下の例を見てみよう。

文字列（4つ）：abc　acb　bac　cab
アナグラム数：4
アナグラムの組：
(abc, acb), (abc, bac), (abc, cab), (acb, bac), (acb, cab), (bac, cab)
アナグラムの組の数：
3 + 2 + 1 = 6

ほら、求められているでしょう。

というわけで、この計算式を用いて作っていこう。
実際に擬似的なコードにしてみると以下のようになる。

// 入力（文字列）を受け取り、ソートしてアナグラム数をカウント
for i := 0; i < n; i++ {
    t := scan()
    strs[sorted(t)]++
}

// アナグラム数からアナグラムの組の数を求め、合計していく
var count int
for _, anagram_num := range strs {
    // nC2の計算部分
    for i := 1; i < anagram_num; i++ {
        count += i
    }
}

ソートしてループして数えて終わり！
とてもシンプルに書けた。

あとは、これに入出力処理などを付け加えれば完成だー！

解答

最終的には以下のようなコードに。

package main

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

func main() {
    var n int
    var t string
    s := wordScanner(100)
    scanInts(s, &n)
    l := make(map[string]int, n)
    for i := 0; i < n; i++ {
    	scanStrings(s, &t)
	l[sorted(t)]++
    }

    var c int
    for _, val := range l {
    	for i := 1; i < val; i++ {
            c += i
        }
    }

    fmt.Println(c)
}

type RuneSlice []rune

func (p RuneSlice) Len() int {
    return len(p)
}

func (p RuneSlice) Less(i, j int) bool {
    return p[i] < p[j]
}

func (p RuneSlice) Swap(i, j int) {
    p[i], p[j] = p[j], p[i]
}

func sorted(s string) string {
    runes := []rune(s)
    sort.Sort(RuneSlice(runes))
    return string(runes)
}

func wordScanner(n int) *bufio.Scanner {
    s := bufio.NewScanner(os.Stdin)
    s.Split(bufio.ScanWords)
    b := make([]byte, n)
    s.Buffer(b, n)
    return s
}

func scanInts(s *bufio.Scanner, vals ...*int) {
    for i := range vals {
	s.Scan()
        n, _ := strconv.Atoi(s.Text())
    	*vals[i] = n
    }
}

func scanStrings(s *bufio.Scanner, vals ...*string) {
    for i := range vals {
        s.Scan()
        *vals[i] = s.Text()
    }
}

解答はシンプルだったけども、実は、コンテスト中に一度この問題に詰まってしまっていた...
計算量の見積もりを間違えて、時間内に実行が終わらない！？なんておバカな勘違いをし、時間を無駄にかけてしまった。

しっかり問題文を読んで、計算量の見積もりを間違えないようにしないと。せっかく解けるのにもったいないしね。

さぁ次回は何を解こうかな。