std::priority_queue を使って複数の整列済みデータをマージ

メモ代わりに,sort -m と同じような動作をするプログラムのソースコードを貼り付けておきます.(行単位で)整列済みのファイルを複数マージして,一つの整列済みファイルを作成するためのプログラムです.入力の指定にはコマンドライン引数,出力には標準出力(std::cout)を利用しています.

入力ストリームのポインタを優先順序付きキュー(std::priority_queue)に詰め直すという処理を繰り返しているだけです.効率を重視する場合,std::string や std::getline() の利用は避けて,自分でバッファリングする方がよいと思います.また,Node の代わりに std::pair と std::greater<略> を使うというのもアリでしょう.

基本的には sort -m を使えばよいのですが,入力が gzip や bzip2 で圧縮されている場合に利用できるほか,処理単位が「行」でない場合など,いろいろな状況で使えます.std::priority_queue は便利です.

#include <fstream>
#include <iostream>
#include <queue>
#include <string>
#include <vector>

#include <boost/shared_ptr.hpp>

namespace {

class Merger
{
public:
	Merger() : queue_() {}

	// Adds an input stream.
	bool Add(std::istream *input)
	{
		Node node(input);
		if (!node.Read())
			return false;

		queue_.push(node);
		return true;
	}

	// Reads the next line.
	bool Read(std::string *line)
	{
		if (queue_.empty())
			return false;

		const Node &top_node = queue_.top();
		std::istream *input = top_node.input();
		*line = top_node.line();
		queue_.pop();

		Add(input);
		return true;
	}

private:
	class Node
	{
	public:
		Node() : input_(NULL), line_() {}
		explicit Node(std::istream *input) : input_(input), line_() {}

		// Reads the next line.
		bool Read() { return std::getline(*input_, line_); }

		// Accesses member variables.
		std::istream *input() const { return input_; }
		const std::string &line() const { return line_; }

		// Comparison function for priority queue.
		class Comparer
		{
		public:
			bool operator()(const Node &lhs, const Node &rhs) const
			{
				return lhs.line_ > rhs.line_;
			}
		};

	private:
		std::istream *input_;
		std::string line_;
	};

private:
	std::priority_queue<Node, std::vector<Node>, Node::Comparer > queue_;

	// Disallows copies.
	Merger(const Merger &);
	Merger &operator=(const Merger &);
};

}  // namespace

int main(int argc, char *argv[])
{
	Merger merger;

	// Opens input files and makes a list of file streams.
	std::vector<boost::shared_ptr<std::ifstream> > file_list;
	for (int i = 1; i < argc; ++i)
	{
		boost::shared_ptr<std::ifstream> file(new std::ifstream(
			argv[i], std::ios_base::in | std::ios_base::binary));
		if (!*file)
		{
			std::cerr << "error: failed to open file: "
				<< argv[i] << std::endl;
			return 1;
		}
		file_list.push_back(file);

		// Adds file streams to the merger.
		merger.Add(file.get());
	}

	// Reads lines in merged order.
	std::string line;
	while (merger.Read(&line))
		std::cout << line << '\n';

	return 0;
}