ABC313 Dの逆行列解法とライブラリ化【解説・ライブラリ】

ABC313のD問題について、逆行列を用いた解法を思いついたので紹介します。ついでにライブラリ化もしました。

書いてる人について

atcoderで競プロしています。青タッチしてますが、現在は水コーダーです。連続で冷えててつらい

扱う問題

本番ではACできませんでしたが、考え方そのものは間違ってなかったです。

注意
以下この問題のネタバレを含みます。また、とても初歩的ですがある程度の線形代数の知識を要求しています。(逆行列が分かれば読めると思います)

問題概要

インタラクティブな問題になっています。$n$個の01列がどこかに存在し、そのうち 奇数$k$個の要素の和の 偶奇 を最大$n$回聞くことができます。
すべての要素を決定し、答えろという問題です。

線形代数的な考察

以下ではすべての数をmod 2で考えることにします。
$i$回目に聞く質問に$j$番目の要素が含まれる場合に$(i,\space j)$の位置を1、そうでない場合0とする$n\times n$行列$M$を考えます。
また$i$番目の要素を$i$回目の質問の答えとした縦ベクトルを$\boldsymbol b$とし、求める配列をそのまま縦ベクトルとしたものを$\boldsymbol a$としましょう。
このとき、$M \cdot \boldsymbol a = \boldsymbol b$となります。(以下に$n = 5,\space k = 3$の例を示します)

\begin{pmatrix}
    1 & 1 & 1 & 0 & 0\\
    0 & 1 & 1 & 1 & 0\\
    0 & 0 & 1 & 1 & 1\\
    1 & 0 & 0 & 1 & 1\\
    1 & 1 & 0 & 0 & 1\\
\end{pmatrix}
\cdot
\begin{pmatrix}
    1\\
    0\\
    0\\
    1\\
    1\\
\end{pmatrix}
=
\begin{pmatrix}
    1\\
    1\\
    0\\
    1\\
    0\\
\end{pmatrix}

この例では、[1, 0, 0, 1, 1]が答えるべき配列であり、例えば3回目の質問では[3, 4, 5]と質問して、0を回答として受けとったことを示しています。

ここで、$M$の逆行列$M^{-1}$が求まるならば、$\boldsymbol a = M^{-1} \cdot \boldsymbol b$として容易に$\boldsymbol a$を求めることができます。しかし、解決すべき問題として 逆行列が存在するのか ということと、計算量が間に合うのか ということがあります。

逆行列の存在

逆行列が存在するためには、各質問に相当する横ベクトルが線形独立であることが必要十分です。
これは容易に構築できます。

例えば、$(0 \space 1 \space 1\space 1)$ と $(1\space 0\space 1\space 1)$は線形独立です。このように、0の位置をずらしていくことで$k+1$個の線形独立なベクトルを得ることができます。残りの$(n - k -1 )$個は斜めに並べていけばいいです。

\begin{pmatrix}
    0 & 1 & 1 & 1 & 1 & 1 & 0 & 0 & 0\\
    1 & 0 & 1 & 1 & 1 & 1 & 0 & 0 & 0\\
    1 & 1 & 0 & 1 & 1 & 1 & 0 & 0 & 0\\
    1 & 1 & 1 & 0 & 1 & 1 & 0 & 0 & 0\\
    1 & 1 & 1 & 1 & 0 & 1 & 0 & 0 & 0\\
    1 & 1 & 1 & 1 & 1 & 0 & 0 & 0 & 0\\
    0 & 0 & 1 & 1 & 1 & 1 & 1 & 0 & 0\\
    0 & 0 & 0 & 1 & 1 & 1 & 1 & 1 & 0\\
    0 & 0 & 0 & 0 & 1 & 1 & 1 & 1 & 1\\
\end{pmatrix}

$n = 9,\space k = 5$のとき、上のようになります。 いろんな構築方法があると思われるので、工夫してみてください。

計算量

このアルゴリズムでは逆行列を求める部分がボトルネックとなります。$n$は最大で1000であり、逆行列を求めるのには$O(n ^3)$かかるため、工夫なしでは厳しいです。(内部の演算をすべてbit演算に置き換えれば、高速な言語はそのままでも通せるかもしれません)

しかし、横ベクトルをbitsetで持ち、掃き出し法をbit演算で記述することで、かなり高速化することができます。(いわゆる64倍bit高速化)

この工夫により、十分高速に求めることができます。計算量は、ワードサイズを$w$として、$O(n^3 /w)$となります。

ライブラリ化

競プロ文脈では行列はmod 2の文脈で用いられることが多い気がします。そこで、bitsetを用いたものをライブラリ化しておくことにしました。こんなんいつ使うねん

#include <bits/stdc++.h>
using namespace std;

struct bit_square_matrix {
    int n;
    static const int N = 1000;
    static const int M = N * 2;
    using V = bitset<N>;
    vector<V> d;
    bit_square_matrix() : n(0) {}
    bit_square_matrix(int n) : n(n), d(n) {}
    const V &operator[](int i) const {
        return d[i];
    }
    V &operator[](int i) {
        return d[i];
    }

    bit_square_matrix &unit() {
        *this = bit_square_matrix(n);
        for(int i = 0; i < n; i++) d[i][i] = 1;
        return *this;
    }

    bit_square_matrix operator+(const bit_square_matrix &a) const {
        assert(n == a.n);
        bit_square_matrix r(n);
        for(int i = 0; i < n; i++)
            for(int j = 0; j < n; j++) {
                r[i][j] = d[i][j] ^ a[i][j];
            }
        return r;
    }

    bit_square_matrix &operator+=(const bit_square_matrix &B) {
        for(int i = 0; i < n; i++) {
            for(int j = 0; j < n; j++) {
                (*this)[i][j] = (*this)[i][j] ^ B[i][j];
            }
        }
        return *this;
    }

    bit_square_matrix operator*(const bit_square_matrix &a) const {
        assert(n == a.n);
        bit_square_matrix r(n);

        for(int i = 0; i < n; i++) {
            for(int k = 0; k < n; k++) {
                for(int j = 0; j < n; j++) {
                    r[i][j] = (d[i][k] & a[k][j]) ^ r[i][j];
                }
            }
        }
        return r;
    }

    bit_square_matrix &operator*=(const bit_square_matrix &a) {
        assert(n == a.n);
        vector<V> r(n);
        for(int i = 0; i < n; i++) {
            for(int k = 0; k < n; k++) {
                for(int j = 0; j < n; j++) {
                    r[i][j] = ((*this)[i][k] & a[k][j]) ^ r[i][j];
                }
            }
        }
        d.swap(r);
        return (*this);
    }

    vector<int> operator*(const vector<int> &B) const {
        vector<int> C(n);
        assert((int)B.size() == n);
        for(int i = 0; i < n; i++) {
            for(int j = 0; j < n; j++) {
                C[i] ^= (*this)[i][j] & B[j];
            }
        }
        return C;
    }

    bool operator==(const bit_square_matrix &b) {
        if(n != b.n) return false;
        for(int i = 0; i < n; i++) {
            if((*this)[i] != b[i]) return false;
        }
        return true;
    }

    bit_square_matrix pow(long long t) const {
        if(!t) return bit_square_matrix(n).unit();
        if(t == 1) return *this;
        bit_square_matrix r = pow(t >> 1);
        r = r * r;
        if(t & 1) r = r * (*this);
        return r;
    }

    // 逆行列を返す。掃き出し法。存在しない場合は長さ0の配列を返す。
    bit_square_matrix Inv() {
        vector<bitset<M>> B(n);
        for(int i = 0; i < n; i++) {
            for(int j = 0; j < n; j++) {
                B[i][j] = d[i][j];
            }
            B[i][i + n] = 1;
        }

        for(int i = 0; i < n; i++) {
            // 注目している変数の係数の絶対値が大きい式をi番目に持ってくる
            int pivot = i;
            for(int j = i; j < n; j++) {
                if(B[j][i] != 0) {
                    pivot = j;
                    break;
                }
            }
            swap(B[i], B[pivot]);

            // 逆行列存在しない
            if(B[i][i] == 0) return bit_square_matrix();

            for(int j = 0; j < n; j++) {
                if(i != j) {
                    if(B[j][i] == 0) continue;
                    B[j] ^= B[i];
                }
            }
        }
        bit_square_matrix res(n);
        for(int i = 0; i < n; i++) {
            for(int j = 0; j < n; j++) res[i][j] = B[i][j + n];
        }
        return res;
    }

    // Ax=bとなるx。Aの逆行列が存在しない場合、空の配列を返す。
    vector<int> gauss_jordan(vector<int> b) {
        bit_square_matrix inv = (*this).Inv();
        if(inv.n == 0) return {};
        b = inv * b;
        return b;
    }

    friend ostream &operator<<(ostream &os, bit_square_matrix &p) {
        for(int i = 0; i < p.n; i++) {
            for(int j = 0; j < p.n; j++) {
                os << p[i][j] << " ";
            }
            os << endl;
        }
        return (os);
    }
};

これを貼って通すことができました。(150 ms)

コードの正確性は保証しかねますが、あくまで競プロで使いたい方は使ってみてください。

感想など

コンテスト中には線形独立なベクトルの構築まではできていましたが、$n$の制約を100までだと思い込んでしまい、TLEの理由がわからず通せませんでした。ちょっと考えればbitset高速化は思いつけたと思うので、とてもくやしいですが、またがんばって精進していこうと思います。

最後に

ご意見ご感想などあればぜひお知らせください。応援メッセージでもいいのよ

ではでは