samber/loを使ってABC329のA-Fを解いてみた

はじめに

これは、富士通クラウドテクノロジーズ Advent Calendar 2023 の5日目の記事です。

投稿動機

その1

競技プログラミングのコツの一つとして、ライブラリを利用してコードの記述量を減らし、より早く問題を解くというものがあります。競技プログラミング用の自作ライブラリを設計する際は、この観点をはじめとして、使い勝手の良さ、コマンドの覚えやすさなどの様々な観点を加味しながら設計を行っていきます。他作のライブラリを使ってみたり、汎用のライブラリを触ってみたりもします。そんな中でlodashのエッセンスを取り込めるsamber/loというライブラリを見つけたので、これを利用した場合にどのようなコードになるかを検証してみました。

その2

プログラミング文体練習というオライリー本を読んで、プログラミングの文体というものは、使用言語に依らず書き手が幅広い選択肢の中から選ぶものである、ということを学んだためGoを用いて様々な文法でコードを書いてみたいと思いました。

対象読者

• Goで競技プログラミングをしてみたい方
• 言語問わず、自作ライブラリの使い勝手に悩んでいる方

samber/lo

samber/loとは、lodashスタイルでGoを記述できるGenerics対応のライブラリになります。Sliceに対してはFilterやMap、FlatMapなどをはじめとして約40種類の関数を利用でき、Mapに対しては、Keys, Values, MapToSliceなどの関数を利用できるほか、Tuple型の実装や、Zip関数なども利用可能です。
また今回の利用はありませんでしたが、非同期実行、エラーハンドリング、リトライ処理などの関数も利用できます。

samber/loは2023年12月時点でAtCoderの提出先環境に用意されていません。samber/loなどの外部ライブラリを利用したコードを提出する際は、依存関係を洗い出して1ファイル化する処理などが必要になります。
主なライブラリとしてgottaniが挙げられますが、更新頻度が低く、ジェネリクスの構文木におそらく非対応のためsamber/loとは相性が悪いです（go.mod内のgoのversionを上げて、依存ライブラリのバージョンを上げて自前ビルドすれば少しマシになります）。

解いたコンテスト

今回はＳｋｙ株式会社プログラミングコンテスト2023（AtCoder Beginner Contest 329）のA-F問題を解きました。

Aから順に、問題を解く際に利用したsamber/loの機能をについての紹介と、利用した感想について書いていきたいと思います。

A問題

文字列$S$を空白で区切って出力せよ

コード

package main

import (
        "bufio"
        "fmt"
        "math"
        "os"

        "github.com/samber/lo"
)

var (
        In  = bufio.NewScanner(os.Stdin)
        Out = bufio.NewWriter(os.Stdout)
        Dbg = bufio.NewWriter(os.Stderr)
)

func init() {
        In.Split(bufio.ScanWords)
        In.Buffer([]byte{}, math.MaxInt64)
}

func Reads() string {
        In.Scan()
        return In.Text()
}

type Input struct {
        s string
}

type Output struct {
        ans string
}

func IOInput() Input {
        s := Reads()
        return Input{
                s: s,
        }
}

func Solve(in Input) Output {
        ans := string(lo.FlatMap([]rune(in.s), func(item rune, index int) []rune {
                return lo.If(index == 0, []rune{item}).Else([]rune{' ', item})
        }))
        return Output{ans: ans}
}

func IOOutput(ans Output) {
        fmt.Fprintln(Out, ans.ans)
}

func main() {
        defer Out.Flush()
        input := IOInput()
        ans := Solve(input)
        IOOutput(ans)
}

感想

普通に解く場合は

ans := strings.Jonin(strings.Split(s, ""), " ")

のような形で問題なさそうですが、今回は検証もかねてFlatMapとIf関数を利用しました。
FlatMapとは、機能的にMapの上位互換のようなもので何かと便利です。返り値のSliceの要素数を制御できます。今回のコードでも、与えられた[]runeに対して、$\mathrm{index}=1$以外の要素では$空白+文字$を返すことで、返り値に空白区切りの文字列を得ることができます。

B問題

数列$A_1, ... A_N$の中で二番目に大きな数を出力せよ。

コード

package main

import (
        "bufio"
        "fmt"
        "math"
        "os"
        "strconv"

        "github.com/samber/lo"
)

var (
        In  = bufio.NewScanner(os.Stdin)
        Out = bufio.NewWriter(os.Stdout)
        Dbg = bufio.NewWriter(os.Stderr)
)

func init() {
        In.Split(bufio.ScanWords)
        In.Buffer([]byte{}, math.MaxInt64)
}

func Reads() string {
        In.Scan()
        return In.Text()
}

func Readi() int {
        return lo.Must(strconv.Atoi(Reads()))
}

type Input struct {
        n int
        a []int
}

type Output struct {
        ans int
}

func IOInput() Input {
        n := Readi()
        a := lo.Map(lo.Range(n), func(item int, index int) int {
                return Readi()
        })
        return Input{
                n: n,
                a: a,
        }
}

func Solve(in Input) Output {
        aMax := lo.Max(in.a)
        ans := lo.Max(lo.Filter(in.a, func(item int, index int) bool { return item != aMax }))
        return Output{ans: ans}
}

func IOOutput(ans Output) {
        fmt.Fprintln(Out, ans.ans)
}

func main() {
        defer Out.Flush()
        input := IOInput()
        ans := Solve(input)
        IOOutput(ans)
}

感想

配列$a$を受け取る際に、Map関数とRange関数を用いました。
Range関数はPythonのrange表記と使い勝手が似ていますが、引数を一つしかとれません。

result := lo.Range(4)
// [0, 1, 2, 3]

result := lo.RangeFrom(1, 5)
// [1, 2, 3, 4, 5]

result := lo.RangeWithSteps(0, 20, 5)
// [0, 5, 10, 15]

https://github.com/samber/lo#range--rangefrom--rangewithsteps

また、start, endや、start, end, stepsを引数にとれる関数もありますが、関数名が長いので、おとなしく普通のfor文を書いたほうが見た目がスッキリします。

C問題

$S$の部分文字列のうち、文字種一種類からなるものの数を求めよ。

コード

package main

import (
        "bufio"
        "fmt"
        "math"
        "os"
        "strconv"

        "github.com/samber/lo"
)

var (
        In  = bufio.NewScanner(os.Stdin)
        Out = bufio.NewWriter(os.Stdout)
        Dbg = bufio.NewWriter(os.Stderr)
)

func init() {
        In.Split(bufio.ScanWords)
        In.Buffer([]byte{}, math.MaxInt64)
}

func Reads() string {
        In.Scan()
        return In.Text()
}

func Readi() int {
        return lo.Must(strconv.Atoi(Reads()))
}

func CharIdx() map[rune]int {
        return lo.Associate(lo.LowerCaseLettersCharset, func(r rune) (rune, int) {
                return r, 0
        })
}

type Input struct {
        n int
        s string
}

type Output struct {
        ans int
}

func IOInput() Input {
        n := Readi()
        s := Reads()
        return Input{
                n: n,
                s: s,
        }
}

func Solve(in Input) Output {
        counts := CharIdx()
        l := 0
        r := 0
        for l < in.n && r < in.n {
                if in.s[l] == in.s[r] {
                        counts[[]rune(in.s)[l]] = lo.Max([]int{counts[[]rune(in.s)[l]], r - l + 1})
                        r++
                } else {
                        l = r
                }
        }
        return Output{ans: lo.Sum(lo.Values(counts))}
}

func IOOutput(ans Output) {
        fmt.Fprintln(Out, ans.ans)
}

func main() {
        defer Out.Flush()
        input := IOInput()
        ans := Solve(input)
        IOOutput(ans)
}

感想

func Readi() int {
        return lo.Must(strconv.Atoi(Reads()))
}

Must関数を用いることで、value, err := Hoge()のerrorハンドリングをもみ消すことができます。が、コンテスト中はvalue, _ = Hoge()と書くのと大差なさそうなので、主に自作ライブラリ内で利用するとよさそうです。

func CharIdx() map[rune]int {
        return lo.Associate(lo.LowerCaseLettersCharset, func(r rune) (rune, int) {
                return r, 0
        })
}
// map[rune]int{'a': 0, 'b':0, ..., 'y': 0, 'z': 0}

lo.LowerCaseLettersCharsetは英小文字の[]runeです。ほかにも英大文字、大文字と小文字と数字、などいろいろなCharsetがあります。

var (
	LowerCaseLettersCharset = []rune("abcdefghijklmnopqrstuvwxyz")
	UpperCaseLettersCharset = []rune("ABCDEFGHIJKLMNOPQRSTUVWXYZ")
	LettersCharset          = append(LowerCaseLettersCharset, UpperCaseLettersCharset...)
	NumbersCharset          = []rune("0123456789")
	AlphanumericCharset     = append(LettersCharset, NumbersCharset...)
	SpecialCharset          = []rune("!@#$%^&*()_+-=[]{}|;':\",./<>?")
	AllCharset              = append(AlphanumericCharset, SpecialCharset...)
)

https://pkg.go.dev/github.com/samber/lo#pkg-variables

lo.Associate関数はSliceからMapへの変換を表現できます。今回の使用用途としては、SliceのruneであるLowerCaseLettersCharsetをの各値をkey、0をvalueとしてmap[rune]intを作成しました。

func Solve(in Input) Output {
        counts := CharIdx()
        l := 0
        r := 0
        for l < in.n && r < in.n {
                if in.s[l] == in.s[r] {
                        counts[[]rune(in.s)[l]] = lo.Max([]int{counts[[]rune(in.s)[l]], r - l + 1})
                        r++
                } else {
                        l = r
                }
        }
        return Output{ans: lo.Sum(lo.Values(counts))}
}

lo.Max, lo.Min関数は引数が[]constraints.Ordered型なので、int, intの比較の際にわざわざスライス型にキャストしてやる必要があり、競プロ目的としてはやや使い勝手が悪い印象でした。

func Max[T constraints.Ordered](values ...T) T {
    return lo.Max(values)
}

自分は上のようにラップして利用しようと思いました。

D問題

人$A_1,...,A_N$の中から一人を選ぶ投票に$M$票の投票があった。
$i$番目の票は$A_i$へ投票されている。
$1 \le i \le N$において、$i$票目まで開票されているときのその時点での当選候補者を出力せよ。

コード

package main

import (
        "bufio"
        "fmt"
        "math"
        "os"
        "strconv"

        "github.com/samber/lo"
)

var (
        In  = bufio.NewScanner(os.Stdin)
        Out = bufio.NewWriter(os.Stdout)
        Dbg = bufio.NewWriter(os.Stderr)
)

func init() {
        In.Split(bufio.ScanWords)
        In.Buffer([]byte{}, math.MaxInt64)
}

func Reads() string {
        In.Scan()
        return In.Text()
}

func Readi() int {
        return lo.Must(strconv.Atoi(Reads()))
}

type Input struct {
        n int
        m int
        a []int
}

type Output struct {
        ans []int
}

func IOInput() Input {
        n := Readi()
        m := Readi()
        a := lo.Map(lo.Range(m), func(item int, index int) int { return Readi() - 1 })
        return Input{
                n: n,
                m: m,
                a: a,
        }
}

func Solve(in Input) Output {
        ans := make([]int, in.m)
        counts := make([]int, in.n)
        iTop := 0
        for i, ia := range in.a {
                counts[ia]++
                iTop = lo.If(counts[ia] > counts[iTop], ia).
                        ElseIf(counts[ia] == counts[iTop], lo.Min([]int{iTop, ia})).
                        Else(iTop)
                ans[i] = iTop + 1
        }
        return Output{ans: ans}
}

func IOOutput(ans Output) {
        for i := range ans.ans {
                fmt.Fprintln(Out, ans.ans[i])
        }
}

func main() {
        defer Out.Flush()
        input := IOInput()
        ans := Solve(input)
        IOOutput(ans)
}

感想

func Solve(in Input) Output {
        ans := make([]int, in.m)
        counts := make([]int, in.n)
        iTop := 0
        for i, ia := range in.a {
                counts[ia]++
                iTop = lo.If(counts[ia] > counts[iTop], ia).
                        ElseIf(counts[ia] == counts[iTop], lo.Min([]int{iTop, ia})).
                        Else(iTop)
                ans[i] = iTop + 1
        }
        return Output{ans: ans}
}

lo.If関数はチェーンでlo.ElseIfをつなげて書くことができ、最後にElseを書くことで完結するものです。
それぞれの第一引数の条件を満たす場合は第二引数の値が返るようになります。今回のlo.If関数をライブラリを使わずに書くと以下のようになります。

if counts[ia] > counts[iTop] {
        iTop = ia
} else if counts[ia] == counts[iTop] {
        iTop = lo.Min([]int{iTop, ia})
}

lo.If関数の場合は、Ifの条件に依らず必ずiTopに何らかの値が代入されるということを確認しやすいですが、単純なif文で書いたほうが見慣れていると感じました。

E問題

文字列$S, T, X$が与えられるので、$X$の中の部分文字を$T$に置き換えることを繰り返して$S$となるか判定せよ

コード

package main

import (
        "bufio"
        "fmt"
        "math"
        "os"
        "strconv"

        "github.com/samber/lo"
        "golang.org/x/exp/constraints"
)

var (
        In  = bufio.NewScanner(os.Stdin)
        Out = bufio.NewWriter(os.Stdout)
        Dbg = bufio.NewWriter(os.Stderr)
)

func init() {
        In.Split(bufio.ScanWords)
        In.Buffer([]byte{}, math.MaxInt64)
}

func Reads() string {
        In.Scan()
        return In.Text()
}

func Readi() int {
        return lo.Must(strconv.Atoi(Reads()))
}

type Input struct {
        n int
        m int
        s []rune
        t []rune
}

type Output struct {
        ans bool
}

func IOInput() Input {
        n := Readi()
        m := Readi()
        s := []rune(Reads())
        t := []rune(Reads())
        return Input{
                n: n,
                m: m,
                s: s,
                t: t,
        }
}

func All[T any](vals []T, f func(i int, v T) bool) bool {
        for i, v := range vals {
                if !f(i, v) {
                        return false
                }
        }
        return true
}

func isGoodI(s, t []rune) bool {
        return All(s, func(i int, v rune) bool {
                return v == '#' || v == t[i]
        })
}

func Max[T constraints.Ordered](vals ...T) T {
        return lo.Max(vals)
}

func Min[T constraints.Ordered](vals ...T) T {
        return lo.Min(vals)
}

func Solve(in Input) Output {
        goodI := lo.FlatMap(in.s[:in.n-in.m+1], func(r rune, i int) []int {
                return lo.If(isGoodI(in.s[i:i+in.m], in.t), []int{i}).Else([]int{})
        })
        visited := make([]bool, in.n)
        for len(goodI) > 0 {
                i := goodI[0]
                goodI = goodI[1:]
                visited[i] = true
                for j := 0; j < in.m; j++ {
                        in.s[i+j] = '#'
                }
                l, r := Max(0, i-in.m+1), Min(i+in.m, in.n-in.m+1)
                for j := l; j < r; j++ {
                        if isGoodI(in.s[j:j+in.m], in.t) && !visited[j] {
                                goodI = append(goodI, j)
                        }
                }
        }
        ans := All(in.s, func(i int, v rune) bool { return v == '#' })
        return Output{ans: ans}
}

func IOOutput(ans Output) {
        fmt.Fprintln(Out, lo.If(ans.ans, "Yes").Else("No"))
}

func main() {
        defer Out.Flush()
        input := IOInput()
        ans := Solve(input)
        IOOutput(ans)
}

感想

func All[T any](vals []T, f func(i int, v T) bool) bool {
        for i, v := range vals {
                if !f(i, v) {
                        return false
                }
        }
        return true
}

loのAll関数（EveryBy関数）では、第二引数に与える関数にindexの情報を与えられず、今回のケースで利用できなかったので自作しました。
また、今回はSolve関数の中身にうまくloの関数をあてることができませんでした。このあたりの難易度からは高速化のために状態を持つことが増えていくため、loの関数が使いづらくなってきます

F問題

$N$個の箱に色のついたボールが入っている。
$Q$個のクエリが与えられる。
クエリ: 箱aから箱bに中のボールをすべて移し、bの中に何種類の色のボールが入っているかを標準出力せよ

コード

package main

import (
        "bufio"
        "fmt"
        "math"
        "os"
        "strconv"

        "github.com/samber/lo"
)

var (
        In  = bufio.NewScanner(os.Stdin)
        Out = bufio.NewWriter(os.Stdout)
        Dbg = bufio.NewWriter(os.Stderr)
)

func init() {
        In.Split(bufio.ScanWords)
        In.Buffer([]byte{}, math.MaxInt64)
}

func Reads() string {
        In.Scan()
        return In.Text()
}

func Readi() int {
        return lo.Must(strconv.Atoi(Reads()))
}

type Input struct {
        n  int
        nq int
        c  []int
        qa []int
        qb []int
}

type Output struct {
        ans []int
}

func IOInput() Input {
        n, nq := Readi(), Readi()
        c := lo.Map(lo.Range(n), func(item int, index int) int { return Readi() - 1 })
        qa, qb := make([]int, nq), make([]int, nq)
        for i := range lo.Range(nq) {
                qa[i], qb[i] = Readi()-1, Readi()-1
        }
        return Input{
                n:  n,
                nq: nq,
                c:  c,
                qa: qa,
                qb: qb,
        }
}

func Solve(in Input) Output {
        boxes := lo.Map(in.c, func(v, i int) map[int]any { return map[int]any{v: nil} })
        ans := lo.Map(lo.Zip2(in.qa, in.qb), func(q lo.Tuple2[int, int], i int) int {
                na, nb := len(boxes[q.A]), len(boxes[q.B])
                if na <= nb {
                        for j := range boxes[q.A] {
                                boxes[q.B][j] = nil
                        }
                } else {
                        for j := range boxes[q.B] {
                                boxes[q.A][j] = nil
                        }
                        boxes[q.B] = boxes[q.A]
                }
                boxes[q.A] = make(map[int]any, 0)
                return len(boxes[q.B])
        })

        return Output{ans: ans}
}

func IOOutput(ans Output) {
        lo.ForEach(ans.ans, func(v, i int) { fmt.Fprintln(Out, v) })
}

func main() {
        defer Out.Flush()
        input := IOInput()
        ans := Solve(input)
        IOOutput(ans)
}

感想

func Solve(in Input) Output {
        boxes := lo.Map(in.c, func(v, i int) map[int]any { return map[int]any{v: nil} })
        ans := lo.Map(lo.Zip2(in.qa, in.qb), func(q lo.Tuple2[int, int], i int) int {
                na, nb := len(boxes[q.A]), len(boxes[q.B])
                if na <= nb {
                        for j := range boxes[q.A] {
                                boxes[q.B][j] = nil
                        }
                } else {
                        for j := range boxes[q.B] {
                                boxes[q.A][j] = nil
                        }
                        boxes[q.B] = boxes[q.A]
                }
                boxes[q.A] = make(map[int]any, 0)
                return len(boxes[q.B])
        })

        return Output{ans: ans}
}

lo.Zip2を利用することで、[]int[T]、[]int[U]を、[]Tuple[T,U]に変形でき、lo.Mapで処理できるようにしました。（Pythonのzip関数と同様）

二つのmap[int]any型を一つにまとめる部分は、Assign関数で同様のことができますが、新しいmapを生成して返す実装のため、今回の問題のように要素の多いMapに要素の小さいMapをマージすることで計算量を落とすような問題には利用できませんでした。

まとめ

Mapなどの汎用的なSlice操作、Map操作を行える関数をGoで記述できて、ボイラーテンプレートを書く時のような飽き飽きとした気持ちなくコードを書くことができて気持ち良かったです。使いづらい関数に関しても、自作ライブラリの下地として利用すると便利そうという見方もできました。
samber/moも利用するつもりでしたが、SomeやTask、Eitherなどのモナドはあまり競プロで利用してもうまみがわからなかったため今回は省略しました。
lodashの表記はできるものの、チェーンでつなげていくことはできないです。パイプライニングなどをしたい場合はSliceをラップする構造体を自作する必要がありそうです。。Goは関数を値として扱え、高階関数とも相性いので関数型プログラミングに向いている気がしましたが、そうでもないみたいです。

終わりに

明日は @George22e さんの 『ヴァイオリンの音色のよさはスペクトラムアナライザーで分かるか試してみた』です。芸術領域とITの融合はとても面白そうで興味が持てます！楽しみです！