ふかふかブログ

興味のあることについてゆる~く書きます

GolangでPythonのrandom.sampleを実現する

最近趣味でPythonしか書いていなかったので、素振りとしてGoを勉強しているのですが

任意のグループからa個選ぶという処理で少し詰まったので忘備として書いておきます。

Pythonのrandom.sampleについて

任意のグループからa個選ぶというもの。

引数として渡されたリストは変更されません。 [1, 2, 3]部分

>>> random.sample([1, 2, 3], 2)
[2, 1]
>>> random.sample([1, 2, 3], 2)
[1, 3]
>>> random.sample([1, 2, 3], 2)
[3, 1]

9.6. random — 擬似乱数を生成する — Python 3.6.5 ドキュメント

golangで実現

こんな感じになると思います。

func randomSample(s []string, num int) []string {
    rand.Seed(time.Now().UnixNano())
    var sample []string
    n := make([]string, len(s))
    copy(n, s)

    for i := 0; i < num; i++ {
        index := rand.Intn(len(n))
        sample = append(sample, n[index])
        n = append(n[:index], n[index+1:]...)
    }
    return sample
}

ポイントとしては、n := make([]string, len(s)) として元のスライスを変更しないようにしています。

go始めたばかりで右も左もよくわかってないのでツッコミあればこっそり教えてください。

package main

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

func main() {
    s := []string{"a", "b", "c", "d"}
    fmt.Println(randomSample(s, 2)) // [a b] or [a c] or ...
    fmt.Println(s)                  // [a b c d]
}

func randomSample(s []string, num int) []string {
    rand.Seed(time.Now().UnixNano())
    var sample []string
    n := make([]string, len(s))
    copy(n, s)

    for i := 0; i < num; i++ {
        index := rand.Intn(len(n))
        sample = append(sample, n[index])
        n = append(n[:index], n[index+1:]...)
    }
    return sample
}