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

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

ARC 048 B - AtCoderでじゃんけんを (Go)

irisruneです。本田圭佑選手とゲーセンのじゃんけんゲームはどっちの方が強いんでしょうかと思いました。

atcoder.jp

多分アルゴリズムを考えること自体はそこまで難しくはないと思うのですが、実装はやや重いと思います。

package main

import (
    "fmt"
    "sort"
)

type Person struct {
    i    int
    rate int
    hand int
    win  int
    lose int
    draw int
}

type People []Person

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

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

type ByRateHandI struct {
    People
}

func (b ByRateHandI) Less(i, j int) bool {
    if b.People[i].rate != b.People[j].rate {
        return b.People[i].rate < b.People[j].rate
    }
    if b.People[i].hand != b.People[j].hand {
        return b.People[i].hand < b.People[j].hand
    }
    return b.People[i].i < b.People[i].i
}

type ByI struct {
    People
}

func (b ByI) Less(i, j int) bool {
    return b.People[i].i < b.People[j].i
}

type Group struct {
    cumSum int
    hand1  int
    hand2  int
    hand3  int
}

func main() {
    var n int
    fmt.Scan(&n)

    var people People = make([]Person, n)
    for i := range people {
        fmt.Scan(&people[i].rate, &people[i].hand)
        people[i].i = i
        people[i].win, people[i].lose, people[i].draw = 0, 0, 0
    }

    sort.Sort(ByRateHandI{people})

    groups := []Group{}
    group := Group{0, 0, 0, 0}
    rateNow := 0
    for i, p := range people {
        if p.rate != rateNow {
            group.cumSum = i
            groups = append(groups, group)
            group.hand1, group.hand2, group.hand3 = 0, 0, 0
            rateNow = p.rate
        }
        switch p.hand {
        case 1:
            group.hand1++
        case 2:
            group.hand2++
        default:
            group.hand3++
        }
    }
    group.cumSum = n
    groups = append(groups, group)

    iGroup := 0
    for i, p := range people {
        if i >= groups[iGroup].cumSum {
            iGroup++
        }
        people[i].win += groups[iGroup-1].cumSum
        people[i].lose += n - groups[iGroup].cumSum

        switch p.hand {
        case 1:
            people[i].win += groups[iGroup].hand2
            people[i].lose += groups[iGroup].hand3
            people[i].draw += groups[iGroup].hand1 - 1
        case 2:
            people[i].win += groups[iGroup].hand3
            people[i].lose += groups[iGroup].hand1
            people[i].draw += groups[iGroup].hand2 - 1
        default:
            people[i].win += groups[iGroup].hand1
            people[i].lose += groups[iGroup].hand2
            people[i].draw += groups[iGroup].hand3 - 1
        }
    }
    sort.Sort(ByI{people})
    for _, p := range people {
        fmt.Println(p.win, p.lose, p.draw)
    }
}

とりあえず解き方の大筋を。

  1. 元の並び順を保持しながら参加者をレートでソートする(今回は昇順)。
  2. 参加者が1名以上該当するレート毎に、合計人数の累積和とグー/チョキ/パーそれぞれの人数(累積和ではない)を記録する。
  3. 2.の情報を用いて各参加者毎に対戦結果を計算する。
  4. 参加者を元の並び順でソートし直して出力する。

まず、1.と4.でソートを2回行う必要がある関係で、ただでさえ煩わしいGoのソートなのにソート処理だけでそれなりの行数使っています。1.のソートでレート以外の要素も使ったからというのはあるのですが。

そしてこの問題の一番の肝になると思われる2.の部分についてですが、参加者をレート毎にグループ化することでレートでの勝敗と手での勝敗をそれぞれ計算するための情報を記録しています。まずレートのグループ毎に人数の累積和を取ることによってレートでの勝敗結果が線形時間で計算できます。一方グループ内のグー/チョキ/パーそれぞれの人数を記録しておくことで手での勝敗結果も線形時間で計算できます。

以上のアルゴリズムで問題を解いたわけですが、ソートを行っているため計算量はO(N\log N)と解説より多めにかかってしまっています。制約条件が1\leq N \leq 10^5なので遅い言語だと少し最適化が必要かもしれませんね。参加者1人なんてケースあったの今気付いたんですが大会とは一体。

コードレビュー(と見せかけたアルゴリズム面のレビュー)

  • 上でも少し触れたが、1.でのソート対象はレートだけでよい(Goのパッケージは不安定ソートだが今回影響はないため)。
  • 累積和と言いつつ実は求め方が累積和ではない。
  • グループの配列を作成する部分や参加者毎の勝敗情報を計算する部分の処理がやや煩雑でわかりにくい。

雑記

  • 昔のコンテストの問題で配点が大体100点だったりするのはまだいいのですが、101点の問題はAC難易度かなり高くてつらいですね。
  • 社内でRustが話題に出てたので軽く調べましたが入出力が既によくわからないですね…OCaml触ったことはあるのですが。
  • 俺の勝ち!とかどこかでねじ込もうかと思ったけど流石にはっちゃけすぎだと思ったのでやめました。