min-maxの範囲からn個の整数をランダム抽出する方法(ただし重複は含まない)

概要

• min-maxの範囲の整数から、n個の整数をランダム抽出する
• ただし抽出する整数は重複しないものとする
• テストデータ作成用のサンプルコードです

サンプルコード

randomInt.go

package main

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

func allKeys(m map[int]bool) []int {
    i := 0
    result := make([]int, len(m))
    for key, _ := range m {
        result[i] = key
        i++
    }
    return result
}

func pickup(min int, max int, num int) []int {
    numRange := max - min

    selected := make(map[int]bool)
    rand.Seed(time.Now().UnixNano())
    for counter := 0; counter < num; {
        n := rand.Intn(numRange) + min
        if selected[n] == false {
            selected[n] = true
            counter++
        }
    }
    keys := allKeys(selected)
    // ソートしたくない場合は以下1行をコメントアウト
    sort.Sort(sort.IntSlice(keys))
    return keys
}

func main() {
    start := time.Now()
    // 1から1,000,000の範囲でランダムに100個、整数を抽出する
    results := pickup(1, 1000000, 100)
    // 処理時間
    fmt.Println(time.Since(start))
    // 結果
    fmt.Println(results)
}

サンプルコード(setを使用したVer.)

• 非標準のsetを使用したバージョン
• setを[]intに変換する部分がもっとシンプルに書けないかと思ったのですが、[]interface{}から[]intへのキャストはできないっぽい。

randomIntWithSet.go

package main

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

    "github.com/deckarep/golang-set"
)

func pickup2(min int, max int, num int) []int {
    numRange := max - min
    set := mapset.NewSet()

    rand.Seed(time.Now().UnixNano())
    for set.Cardinality() < num {
        n := rand.Intn(numRange) + min
        set.Add(n)
    }
    selected := set.ToSlice()
    ret := make([]int, len(selected))
    for i, x := range selected {
        ret[i] = x.(int)
    }
    // ソートしたくない場合は以下1行をコメントアウト
    sort.Sort(sort.IntSlice(ret))
    return ret
}

func main() {
    start := time.Now()
    // 1から1,000,000の範囲でランダムに100個、整数を抽出する
    results := pickup2(1, 1000000, 100)
    // 処理時間
    fmt.Println(time.Since(start))
    // 結果
    fmt.Println(results)
}

おまけ) []intを1行ごとにファイル出力

output.go

func main() {
    results := pickup(1, 1000000, 100)
    writeIntList("/var/tmp/sample.csv", results)
}

func failOnError(err error) {
    if err != nil {
        log.Fatal("Error:", err)
    }
}

func writeIntList(filepath string, list []int) {
    f, err := os.OpenFile(filepath, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0600)
    failOnError(err)
    defer f.Close()

    var writer *bufio.Writer
    writer = bufio.NewWriter(f)
    for _, v := range list {
        writer.WriteString(strconv.Itoa(v) + "\n")
    }
    writer.Flush()
}

ボツコード

• 最初ソート済みの結果取得のために以下のコードを書いていましたが、これだとmaxの数だけrand.Intn()が実行され、時間がかかってしまいイマイチでした。素直にソートしたほうが速い。

badSample.go

package main

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

func pickupSort(min int, max int, num int) []int {
    numRange := max - min

    rand.Seed(time.Now().UnixNano())
    counter := 0
    selected := make([]int, num)
    for i := 0; i < numRange; i++ {
        if rand.Intn(numRange-i) < (num - counter) {
            selected[counter] = i + min
            counter++
        }
    }
    return selected
}

func main() {
    start := time.Now()
    // 1から1,000,000の範囲でランダムに100個、整数を抽出する
    results := pickupSort(1, 1000000, 100)
    // 処理時間
    fmt.Println(time.Since(start))
    // 結果
    fmt.Println(results)
}

まとめ

• 珠玉のプログラミング 本質を見抜いたアルゴリズムとデータ構造 の「コラム12　サンプリング問題」の擬似コードを、(少しアレンジして)Goで再現してみました。
• お気づきの点がありましたら、お気軽にコメントおねがいします!