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で再現してみました。
• お気づきの点がありましたら、お気軽にコメントおねがいします!