您现在的位置：程式碼之家 >> 編程語言 > >> 更多程式語言 >> 程式設計綜合問答

Linux Shell/Python 隨機抽取文件指定行數寫入新文件：三種工程化實現

編輯：程式設計綜合問答

需求描述

從大文件中隨機抽取 N 行寫入新文件。環境 Linux，語言不限 (Python/awk/Java)。

方案對比

方案	核心命令/代碼	優點	缺點	適用場景
1. shuf (GNU coreutils, 最簡單)	`shuf -n 100 input.txt > output.txt`	一行命令、C 實現極快、無需腳本、支持大文件 (內存映射)	需 GNU shuf (macOS 需 `brew install coreutils` 用 `gshuf`)	臨時需求、文件 < 10GB、交互式
2. Python (標準庫, 跨平台)	見下方完整代碼	無依賴、邏輯可控、易擴展 (加權/分層抽樣)、Windows/Linux/macOS 通用	超大文件需流式算法 (水塘抽樣)	腳本化、需嵌入業務流程、跨平台
3. awk (經典文本處理)	`awk 'BEGIN{srand()} {lines[NR]=$0} END{for(i=1;i<=n;i++) print lines[int(rand()*NR)+1]}' n=100 input.txt > output.txt`	單行、無外部依賴、流式可行	全量載入數組受內存限制、隨機數質量一般	無 Python 環境、文件較小
4. 水塘抽樣 (超大文件/流式, Python)	見下方代碼	O(1) 內存、單遍掃描、嚴格均勻隨機	代碼稍長	文件 > 內存、流式處理、日誌抽樣

Python 完整實現 (標準隨機抽取)

import random
import sys

def random_lines(input_file, output_file, n, encoding='utf-8'):
    """將文件所有行讀入內存隨機抽取 (文件 < 內存)"""
    with open(input_file, 'r', encoding=encoding) as f:
        lines = f.readlines()
    
    if n >= len(lines):
        selected = lines
    else:
        selected = random.sample(lines, n)  # 不重複抽取
    
    with open(output_file, 'w', encoding=encoding) as f:
        f.writelines(selected)

def reservoir_sampling(input_file, output_file, n, encoding='utf-8'):
    """水塘抽樣：單遍掃描，O(1) 內存，嚴格均勻"""
    import itertools
    reservoir = []
    with open(input_file, 'r', encoding=encoding) as f:
        for i, line in enumerate(f):
            if i < n:
                reservoir.append(line)
            else:
                j = random.randint(0, i)
                if j < n:
                    reservoir[j] = line
    with open(output_file, 'w', encoding=encoding) as f:
        f.writelines(reservoir)

if __name__ == '__main__':
    # 用法: python random_lines.py input.txt output.txt 100 [reservoir]
    _, infile, outfile, n, *args = sys.argv
    n = int(n)
    if 'reservoir' in args:
        reservoir_sampling(infile, outfile, n)
    else:
        random_lines(infile, outfile, n)
    print(f"已抽取 {n} 行寫入 {outfile}")

Java 版本 (生產環境常用)

import java.io.*;
import java.nio.file.*;
import java.util.*;

public class RandomLines {
    public static void reservoirSample(Path in, Path out, int n) throws IOException {
        List reservoir = new ArrayList<>(n);
        try (BufferedReader br = Files.newBufferedReader(in)) {
            String line;
            int i = 0;
            while ((line = br.readLine()) != null) {
                if (i < n) reservoir.add(line);
                else {
                    int j = new Random().nextInt(i + 1);
                    if (j < n) reservoir.set(j, line);
                }
                i++;
            }
        }
        Files.write(out, reservoir, StandardCharsets.UTF_8);
    }
    public static void main(String[] args) throws IOException {
        reservoirSample(Paths.get(args[0]), Paths.get(args[1]), Integer.parseInt(args[2]));
    }
}