Yappli Tech Blog

株式会社ヤプリの開発メンバーによるブログです。最新の技術情報からチーム・働き方に関するテーマまで、日々の熱い想いを持って発信していきます。

【Go】イテレータを用いた二分木探索を使ってクリスマスツリーを描いて遊ぶ

Go

はじめに

アドベントカレンダーでクリスマスの日に順番が回ってきたので、クリスマスツリーを描きます🌲🎅

こんにちは、株式会社ヤプリ サーバーサイドエンジニアの坂本(@mohvuba)です。

なぜクリスマスツリー？

ヤプリでは月に1回程度、Go界隈で有名なtenntennさんをお招きしてGoの勉強会を開催しています。 その中でここ最近、ツリー構造のオブジェクトの二分木探索をイテレータを用いて行うという内容のものがありました。

冒頭で書いたように、私がアドベントカレンダーで25日の担当になり、 12/25といえばクリスマス、クリスマスといえばクリスマスツリー！ ということで、タイトル通りの記事を書くに至りました。

クリスマスツリーを描く？そんなネタ、出尽くしてるんじゃないの？と思われるかもしれません。 君のような勘のいいガキは嫌いだよ

ご明察のとおり、クリスマスツリーの描画自体はすでに行われている方の記事(Goでクリスマスツリーを表示する)を見つけることができました。 そこで今回は、その記事およびリポジトリの内容を参考にしつつ、イテレータを取り込むという点に特化したいと思います。*1

いざ、描画

ベースとなるコード

下記のリポジトリのコードを参考にします。*2

ブログ記事もこちらのリポジトリとともに最後に参考文献として貼っておきます。

github.com

今回はこのxmastreeに対して

  • 描画をイテレータに変更する
  • 飾りのパターンを増やし、各飾りがいくつあるかを二分木探索+イテレータで集計する

という二つのことを行いたいと思います。

そのために、以下のような手を加えます。

  • structTreeを改良する
  • オーナメントなど、他のものも出してみる
  • 描画する箇所でイテレータを使う
  • 飾りの各数をイテレータで集計した値を一緒に出してみる

Tree structを改良する & オーナメントなど、他のものも出してみる

イテレータ化、および絵文字の導入にあたってstructをもう少し細分化します。

  • 元のTreeはどちらかというと描画のためのメタ的な設定値と言えるので、TreeSettingとして残します
  • Leafを定義します。Leafはオブジェクト(string)を持たせるようにし、葉または飾りの絵文字を持ちます
  • Branchを定義します。Branchは複数枚のLeaf([]*Leaf)を持ちます
  • 先ほどのTreeSettingと[]*BranchもつTreeを定義します

上記に則ったtree.goのソースコードがこちらです

【tree.go】

package tree

import (
    "fmt"
    "iter"
    "math/rand"
    "time"
)

type (
    ChristmasDecoration string
    TreeSetting         struct {
        MaxWidth    int
        LeafHeight  int 
        StemWidth   int
        StemHeight  int
        TotalHeight int
    }
    Leaf struct {
        Object ChristmasDecoration
    }
    Branch struct {
        Leaves []*Leaf
    }
    Tree struct {
        *TreeSetting
        Branches []*Branch
    }
)

const (
    Default ChristmasDecoration = "¥"
    Bell    ChristmasDecoration = "🔔"
    Star    ChristmasDecoration = "⭐"
    Candy   ChristmasDecoration = "🍬"
    Snow    ChristmasDecoration = "⛄️"
    Present ChristmasDecoration = "🎁"
    Sox     ChristmasDecoration = "🧦"
)

var leafObjectPattarn = []ChristmasDecoration{Default, Default, Default, Default, Default, Default, Default, Default, Default, Default, Default, Default, Bell, Star, Candy, Snow, Present, Sox}

func NewBranch() *Branch {
    return &Branch{}
}

func NewLeaf() *Leaf {
    return &Leaf{
        Object: getLeafObject(),
    }
}

func NewTree(
    branchHeight int,
    stemWidth int,
    stemHeight int,
    totalHeight int,
) *Tree {
    branches := make([]*Branch, branchHeight)
    for i := 0; i < len(branches); i++ {
        branches[i] = NewBranch()
    }
    return &Tree{
        TreeSetting: &TreeSetting{
            MaxWidth:    branchHeight*2 + 1,
            LeafHeight:  branchHeight,
            StemWidth:   stemWidth,
            StemHeight:  stemHeight,
            TotalHeight: totalHeight,
        },
        Branches: branches,
    }
}

func (b *Branch) SetLeaves(amount int) {
    for i := 0; i < amount; i++ {
        b.Leaves = append(b.Leaves, NewLeaf())
    }
}

func echoString(str string, amount int) []string {
    strs := make([]string, 0, amount)
    for i := 0; i < amount; i++ {
        strs = append(strs, str)
    }
    return strs
}

func getLeafObject() ChristmasDecoration {
    // ランダムにオブジェクトを与える
    rand := rand.New(rand.NewSource(time.Now().UnixNano()))
    i := rand.Intn(len(leafObjectPattarn))
    return leafObjectPattarn[i]
}

描画する箇所でイテレータを使う

先ほどのstructの改変ができたら描画していきます。 描画のため、Treeに以下のメソッドを実装します

  • 設定値を元にイテレータを定義して返すSet()(iter.Seq[ChristmasDecoration], iter.Seq[string])
  • 葉と幹の描画用のイテレータをそれぞれ受け取って葉と幹を一段ずつ描画するDrawLine(iter.Seq[ChristmasDecoration], iter.Seq[string])

【tree.go】

// Set 木の葉, 幹のを描画するためのイテレータを返す
func (t *Tree) Set() (iter.Seq[ChristmasDecoration], iter.Seq[string]) {
    return func(yield func(ChristmasDecoration) bool) {
            // 葉
            for i, branch := range t.Branches {
                leafAmount := i*2 + 1 - (i/5)*4
                branch.SetLeaves(leafAmount) // 葉の数を設定できるこのタイミングで葉にオブジェクトを設定する
                spaceLength := (t.MaxWidth - leafAmount) / 2
                strs := make([]string, 0, spaceLength*2+leafAmount+1)
                strs = append(strs, echoString(" ", spaceLength)...)
                for _, leaf := range branch.Leaves {
                    strs = append(strs, string(leaf.Object))
                }
                strs = append(strs, echoString(" ", spaceLength)...)
                strs = append(strs, "")
                var str string
                for _, s := range strs {
                    str += s
                }
                if ok := yield(ChristmasDecoration(str)); !ok {
                    return
                }
            }
        },
        func(yield func(string) bool) {
            // 幹
            spaceLength := (t.MaxWidth-t.StemWidth-2)/2 + 3
            for i := 0; i < t.StemHeight; i++ {
                strs := make([]string, 0, spaceLength*2+t.StemWidth+3)
                strs = append(strs, echoString(" ", spaceLength)...)
                strs = append(strs, echoString("[", 1)...)
                strs = append(strs, echoString(" ", t.StemWidth)...)
                strs = append(strs, echoString("]", 1)...)
                strs = append(strs, echoString(" ", spaceLength)...)
                strs = append(strs, "")
                var str string
                for _, s := range strs {
                    str += s
                }
                if ok := yield(str); !ok {
                    return
                }
            }
            strs := make([]string, 0, spaceLength*2+t.StemWidth+3)
            strs = append(strs, echoString(" ", spaceLength)...)
            strs = append(strs, echoString("`", t.StemWidth+2)...)
            strs = append(strs, echoString(" ", spaceLength)...)
            strs = append(strs, "")
            var str string
            for _, s := range strs {
                str += s
            }
            if ok := yield(str); !ok {
                return
            }
            if ok := yield("\r\n"); !ok {
                return
            }
            if ok := yield("\r\n"); !ok {
                return
            }
        }
}

// DrawLine 葉,幹を表示する
func (t *Tree) DrawLine(leaves iter.Seq[ChristmasDecoration], stem iter.Seq[string]) {
    fmt.Print("\033[32m")
    for s := range leaves {
        fmt.Println(s)
    }
    fmt.Print("\033[31m")
    for s := range stem {
        fmt.Println(s)
    }
}

元のコードが一段ごとに葉の枚数を後から決定し描画する仕様になっていたため、それに倣う形でBranchがどのようなLeafを持つかはBranchごとの描画処理の際に決定されるようになっています。 *3

では、main.goからこれらを呼び出してみます。

【main.go】 *4

package main

import (
    "advent-calendar/tree"
    "flag"
    "fmt"
    "iter"
    "os"
    "time"
)

func main() {
    var (
        size  = flag.Int("size", 8, "size")
        speed = flag.Int("speed", 100, "speed")
    )
    flag.Usage = usage
    flag.Parse()
    for s := range display(*size, *speed) {
        fmt.Print(s)
    }
}

func usage() {
    usage := ""
    usage += "Usage: xmastree: [OPTION]\n"
    usage += "  -size size\n"
    usage += "       size of the tree. Default is 8. \n"
    usage += "  -speed speed\n"
    usage += "       speed of display tree. Default is 100. \n"
    _, err := fmt.Fprintf(os.Stderr, usage)
    if err != nil {
        fmt.Println(err)
        os.Exit(2)
    }
}

func display(size, speed int) iter.Seq[string] {
    return func(yield func(string) bool) {
        maxHeight := size*2 + 9

        for i := 0; i < size; i++ {
            if i != 0 {
                if ok := yield(fmt.Sprintf("\033[%vA", maxHeight)); !ok {
                    return
                }
            }
            tree := tree.NewTree(i*2+5, i+1, 2, maxHeight)
            leaf, stem := tree.Set()
            tree.DrawLine(leaf, stem)
            time.Sleep(time.Duration(speed) * time.Millisecond)
        }

        fmt.Print("\033[1A")
    }
}

結果

絵文字のサイズ問題やターミナルの環境依存によってやや崩れはありますが、ちゃんとクリスマスツリーですね！

飾りの各数をイテレータで集計した値を一緒に出してみる

一緒に出すという点についてですが、 Branchを処理するタイミングでしか葉の中身や枚数を決定できないという事情から、ツリーを描画しないとそもそも飾りの集計ができないためそのように対応します。

二分木探索を行うために、Treeがもつ線形の[]*Branchを親子構造に変換するところから始めます。 そのために、TreeにメソッドMakeBranchBTree() errorを追加で実装します。(木なのか枝なのかややこしい...)

【tree.go】

// MakeBranchBTree Branchの木構造を作成する
func (t *Tree) MakeBranchBTree() error {
    if len(t.Branches) == 0 {
        return nil
    }
    for i := 0; i < len(t.Branches); i++ {
        if len(t.Branches[i].Children) > 0 {
            return errors.New("すでに子ノードが存在します")
        }
        leftIndex := 2*i + 1
        rightIndex := leftIndex + 1 // 右の子ノードを左の子ノードから計算

        // 両方のインデックスが範囲外なら終了
        if leftIndex >= len(t.Branches) && rightIndex >= len(t.Branches) {
            break
        }

        // 左の子ノードが範囲内の場合に追加
        if leftIndex < len(t.Branches) {
            t.Branches[i].Children = append(t.Branches[i].Children, t.Branches[leftIndex])
        }

        // 右の子ノードが範囲内の場合に追加
        if rightIndex < len(t.Branches) {
            t.Branches[i].Children = append(t.Branches[i].Children, t.Branches[rightIndex])
        }
    }
    return nil
}

続いて、[]*Branchを受け取って二分木探索をするイテレータを返す関数BranchBSearchIter(*Branch)iter.Seq[string]を作成します。 絵文字より文字列の方が集計しやすそうなので、絵文字と文字列をmapを用いて変換したものをyield()します。

【tree.go】

const (
    DefaultKey = "default"
    BellKey    = "bell"
    StarKey    = "star"
    CandyKey   = "candy"
    SnowKey    = "snow"
    PresentKey = "present"
    SoxKey     = "sox"
)

var decorationMap = map[ChristmasDecoration]string{
    Default: DefaultKey,
    Bell:    BellKey,
    Star:    StarKey,
    Candy:   CandyKey,
    Snow:    SnowKey,
    Present: PresentKey,
    Sox:     SoxKey,
}

func BSearchBranchIter(b *Branch) iter.Seq[string] {
    return func(yield func(string) bool) {
        for _, child := range b.Children {
            for _, leaf := range child.Leaves {
                if ok := yield(decorationMap[leaf.Object]); !ok {
                    return
                }
            }
            for s := range BSearchBranchIter(child) {
                if ok := yield(s); !ok {
                    return
                }
            }
        }
    }
}

再帰的に自身を呼び出してその戻り値であるイテレータを回すのがポイントですね。 イテレータの関数定義の中でそのイテレータを返している関数自体を呼び出してその戻り値のイテレータを回している感じです。*5

main.goで呼び出してみます。

【main.go】

package main

import (
    "advent-calendar/tree"
    "errors"
    "flag"
    "fmt"
    "iter"
    "os"
    "time"
)

func main() {
    var (
        size  = flag.Int("size", 8, "size")
        speed = flag.Int("speed", 100, "speed")
    )
    flag.Usage = usage
    flag.Parse()
    for s := range display(*size, *speed) {
        fmt.Print(s)
    }
    for k, v := range countMap {
        fmt.Printf("\n%s: %d\n", k, v)
    }
}

func usage() {
    usage := ""
    usage += "Usage: xmastree: [OPTION]\n"
    usage += "  -size size\n"
    usage += "       size of the tree. Default is 8. \n"
    usage += "  -speed speed\n"
    usage += "       speed of display tree. Default is 100. \n"
    _, err := fmt.Fprintf(os.Stderr, usage)
    if err != nil {
        fmt.Println(err)
        os.Exit(2)
    }
}

func display(size, speed int) iter.Seq[string] {
    return func(yield func(string) bool) {
        maxHeight := size*2 + 9

        for i := 0; i < size; i++ {
            if i != 0 {
                if ok := yield(fmt.Sprintf("\033[%vA", maxHeight)); !ok {
                    return
                }
            }
            tree := tree.NewTree(i*2+5, i+1, 2, maxHeight)
            leaf, stem := tree.Set()
            tree.DrawLine(leaf, stem)
            // 飾りの数を数える(最後のツリーのみ)
            if i == size-1 {
                if err := aggregate(tree); err != nil {
                    yield(err.Error())
                    return
                }
            }
            time.Sleep(time.Duration(speed) * time.Millisecond)
        }

        fmt.Print("\033[1A")
    }
}

var countMap = map[string]int{
    tree.DefaultKey: 0,
    tree.BellKey:    0,
    tree.CandyKey:   0,
    tree.StarKey:    0,
    tree.SnowKey:    0,
    tree.PresentKey: 0,
    tree.SoxKey:     0,
}

func aggregate(t *tree.Tree) error {
    if err := t.MakeBranchBTree(); err != nil {
        return err
    }
    if len(t.Branches) == 0 {
        return errors.New("Branchが存在しません")
    }
    for s := range tree.BSearchBranchIter(t.Branches[0]) {
        switch s {
        case tree.DefaultKey:
            countMap[tree.DefaultKey]++
        case tree.BellKey:
            countMap[tree.BellKey]++
        case tree.CandyKey:
            countMap[tree.CandyKey]++
        case tree.StarKey:
            countMap[tree.StarKey]++
        case tree.SnowKey:
            countMap[tree.SnowKey]++
        case tree.PresentKey:
            countMap[tree.PresentKey]++
        case tree.SoxKey:
            countMap[tree.SoxKey]++
        }
    }
    return nil
}

main.goで扱われるのは[]*Branchを持ったTree型のインスタンスですが、親子構造を持つため集計に使うのは[]*Branchの中の先頭の要素だけである点に注意してください。*6

結果

最後に

というわけで今回はイテレータと再帰処理を使ってクリスマスツリーを出力してみました！ このような遊び心のあるプログラミングを通してGoの新しい機能が学べる機会があるのはすごくいいですね。

ヤプリでは現在エンジニア採用に力を入れております。 本記事についてや、ヤプリについて具体的なお話を聞いてみたいと思っていただいていた場合は、ぜひカジュアル面談にご応募いただければと思います！ 一緒にGoStudyなどを通じて高め合いましょう！

open.talentio.com

参考記事・リポジトリ

github.com qiita.com

*1:ゼロから描画プログラムを書こうとして、やれピラミッド状を描くなら平方数のタイミングで改行せねばだの、中央に配置するために空白ブロックを用意しておくのかだの細かく考えて時間を浪費していたのは内緒

*2:リポジトリの制作者に連絡し、使用許可をいただいております。

*3:もっと綺麗なやり方はあるかもしれないですが、目的はイテレータを使用することなので、ご容赦ください🙇

*4:ツリーの葉の文字を「¥」以外の絵文字も適用するために必要な改変と、イテレータへの改変以外は行っていませんが、一部改変を行っていない箇所などを便宜上コード欄に一緒に載せています。

*5:何を言ってるかわからないw

*6:このあたりは再帰処理を紹介するために無理やり親子構造を作ったことに起因していて、特殊な感じになってしまっていることをお詫び申し上げます🙇