ほげめも: 複数の配列をソートする

TopCoder Logo

だいぶ前の話ですが TopCoder SRM631 Division II Level Two 問題で「同じ要素数 N の整数配列 A, B を A の値をキーにしてソートする」というコードを制限時間内に書けず落としたので、反省の気持ちをこめて記事にします。

こういう列単位の複数配列セットのソートはプログラミングコンテストでは頻出の問題に見えますが、自分は回答言語の C++ に慣れてなかったこともあり、終了時間が迫る中で書いたコードが動かず半分パニックになってしまいました。

ということで、ここでは C++ による正しいコードを何通りか考え、次回に活かしていきたいと思います。

自分で書く

本番中に考えた、ライブラリに頼らず自分でソートしちゃうコード。 O(n²) で遅いし、 A に同じ値が出てくるとバグるし、いろいろとひどいのですが、問題の制約条件下ではこれで正しい答えが得られます。提出したコードは結局ミスってて challenge phase で撃墜されましたが…。

#include <iostream>
#include <vector>
#include <climits>

using namespace std;

void sort2vectors(vector<int> &av, vector<int> &bv)
{
        int n = av.size();
        vector<int> av2, bv2;
        for (int i = 0, max = INT_MIN; i < n; i++) {
                int k = 0, min = INT_MAX;
                for (int j = 0; j < n; j++) {
                        if (max < av[j] && av[j] < min) {
                                k = j;
                                min = av[j];
                        }
                }
                av2.push_back(av[k]);
                bv2.push_back(bv[k]);
                max = av[k];
        }
        av = av2;
        bv = bv2;
}

int main(int argc, char **argv)
{
        int a[] = { 5, 1, -10, 7, 12, 2, 10, 20 };
        int b[] = { 3, 4, 2, 7, 1, 4, 3, 4 };
        vector<int> av(a, a+sizeof(a)/sizeof(int));
        vector<int> bv(b, b+sizeof(b)/sizeof(int));
        sort2vectors(av, bv);
        for (size_t i = 0; i < av.size(); i++)
                cout << av[i] << " " << bv[i] << endl;
        return 0;
}

これを実行すると次のような出力が得られます。

$ g++ sort2vectors0.cpp 
$ ./a.out
-10 2
1 4
2 4
5 3
7 7
10 3
12 1
20 4

以下、関数 sort2vectors() をいまどきの C++ 標準ライブラリを使ってもっとまともに書けないか考えてみます。特に標準アルゴリズム <algorithm> の std::sort() を活用する方針でいきたいと思います。

構造体や std::pair を使う

配列 A B の各要素をメンバとして持つレコード型を作りその配列をソートする。要するに Python の zip() のように各列の値をまとめて参照できるようにしてからソートします。

レコード型としては新しい構造体を定義してもよいのですが面倒なので、 C++ 標準のペアオブジェクト std::pair を使ってみることにします。

typedef pair<int, int> pair_t;

bool comp(const pair_t &a, const pair_t &b) {
        return a.first < b.first;
}

void sort2vectors(vector<int> &av, vector<int> &bv)
{
        int n = av.size();
        vector<pair_t> p(n);
        vector<int> av2(n), bv2(n);
        for (int i = 0; i < n; i++)
                p[i] = make_pair(av[i], bv[i]);
        sort(p.begin(), p.end(), comp);
        for (int i = 0; i < n; i++) {
                av2[i] = p[i].first;
                bv2[i] = p[i].second;
        }
        av = av2;
        bv = bv2;
}

書いてて思いましたが std::pair を使っても新しい構造体を定義してもコード量は大して変わらず面倒くさいですね。またこの方法は配列のコピーが発生するので効率があまりよくないです。

Functor (関数オブジェクト) を使う

配列の無駄なコピーを避けるため、次は各配列要素へのポインタのようなものをソートすることを考えてみます。具体的には A B のインデックス値を格納した配列 { 0, 1, ..., N-1 } を用意し、それで配列 A の値を参照しつつソートする作戦です。

class Comp {
public:
        const vector<int> &_v;
        Comp(const vector<int> &v): _v(v) {}
        bool operator()(int a, int b) { return _v[a] < _v[b]; }
};

void sort2vectors(vector<int> &av, vector<int> &bv)
{
        int n = av.size();
        vector<int> p(n), av2(n), bv2(n);
        for (int i = 0; i < n; i++)
                p[i] = i;
        sort(p.begin(), p.end(), Comp(av));
        for (int i = 0; i < n; i++) {
                av2[i] = av[p[i]];
                bv2[i] = bv[p[i]];
        }
        av = av2;
        bv = bv2;
}

C の sort() では比較関数のポインタしか渡せないため苦労することが多いのですが、 C++ で functor を使えばこの例の配列 A の参照のように任意のコンテキスト情報を渡すことができるので便利です。また std::sort() の中で比較関数のインライン展開により高速化が期待できます。

Lambda (無名関数) を使う

さらに C++11 で導入された lambda を使うと、わざわざ functor を定義しなくとも 1 行で書けます。

void sort2vectors(vector<int> &av, vector<int> &bv)
{
        int n = av.size();
        vector<int> p(n), av2(n), bv2(n);
        iota(p.begin(), p.end(), 0);
        sort(p.begin(), p.end(), [&](int a, int b) { return av[a] < av[b]; });
        for (int i = 0; i < n; i++) {
                av2[i] = av[p[i]];
                bv2[i] = bv[p[i]];
        }
        av = av2;
        bv = bv2;
}

lambda ブロックでは外側のブロックのローカル変数も普通に参照できます。同様に C++11 で新しく追加された std::iota() も便利。

なお C++11 の lambda については lambda 完全解説という記事がたいへん参考になりました。

他の言語で書くと

ちなみに Ruby で書くとこうなると思います。

#!/usr/bin/env ruby
# coding: utf-8

def sort2vectors(av, bv)
  pv = (0...av.size).to_a.sort {|a, b| av[a] <=> av[b] }
  av2 = pv.map {|i| av[i]}
  bv2 = pv.map {|i| bv[i]}
  # Ruby でこういう値の返し方は普通しない
  av.replace av2
  bv.replace bv2
end

av = [ 5, 1, -10, 7, 12, 2, 10, 20 ]
bv = [ 3, 4, 2, 7, 1, 4, 3, 4 ]
sort2vectors(av, bv)
av.size.times do |i|
  puts "#{av[i]} #{bv[i]}"
end

Python ならこうか。

#!/usr/bin/env python
# coding: utf-8

def sort2vectors(av, bv):
    p = sorted(range(0, len(av)), key=lambda x: av[x])
    # Python でこういう値の返し方は普通しない
    av[:] = [av[x] for x in p]
    bv[:] = [bv[x] for x in p]

av = [ 5, 1, -10, 7, 12, 2, 10, 20 ]
bv = [ 3, 4, 2, 7, 1, 4, 3, 4 ]
sort2vectors(av, bv)
for i in range(0, len(av)):
    print av[i], bv[i]

こういう言語に慣れてるとやっぱり C++11 の lambda は非常に便利に感じますね。可能であればどんどん使っていきたいものです。

なお PHP には array_multisort() という専用関数があるそうです。これは 2 つより多くの配列にも適用できます。

<?php
$av = array(5, 1, -10, 7, 12, 2, 10, 20);
$bv = array(3, 4, 2, 7, 1, 4, 3, 4);
array_multisort($av, $bv);
for ($i = 0; $i < count($av); $i++) {
  echo "{$av[$i]} {$bv[$i]}\n";
}
?>

まとめ

制限時間の厳しいプログラミングコンテストではこういうのを調べてるだけで時間をロスして不利になりますので、迷わず書けるようになりたいものです。

2014/09/21 12:23:50 JST