在Go语言中,实现蓄水池抽样算法

那是什么东西

当你想要从大量文本中随机抽取行时的那个东西。

请详细阅读此链接:http://sucrose.hatenablog.com/entry/2014/01/11/004615。

嗯,只是重新写了一遍而已。

代码 (daima)

package main

import (
    "bufio"
    "flag"
    "fmt"
    "log"
    "math/rand"
    "os"
    "time"
)

var (
    seed       = flag.Int64("s", time.Now().UnixNano(), "random seed")
    samplesNum = flag.Int("n", 1000, "samples num")
)

func main() {
    flag.Parse()

    rand.Seed(*seed)
    k := *samplesNum

    reservoir := make([]string, k)
    scanner := bufio.NewScanner(os.Stdin)

    for i := 0; i < k; i++ {
        if !scanner.Scan() {
            log.Printf("input line less than %d", k)
            break
        }
        reservoir[i] = scanner.Text()
    }

    n := k

    for scanner.Scan() {
        n++
        r := rand.Intn(n)  // [0, n)
        if r < k {
            reservoir[r] = scanner.Text()
        }
    }
    if err := scanner.Err(); err != nil {
        log.Fatal(err)
    }

    for _, s := range reservoir {
        if s != "" {
            fmt.Fprintln(os.Stdout, s)
        }
    }
}

请构建

go build reservoir_sampling.go

执行 (shí

$ wc -l data.csv
16357121 data.csv

$ time cat data.csv | ./reservoir_sampling -n 10000 | wc -l
10000

real  0m11.707s
user  0m2.691s
sys 0m10.424s

从大约1600万个选项中提取10000个样本只需要大约10秒钟。
在我的Mac电脑上,使用Python3花费了大约1分钟的时间,算是相当满意。

广告
将在 10 秒后关闭
bannerAds