なんとか通した。結構迷走した。

問題へのリンク

問題概要

0 と 1 と 2 のみからなる $N \times K$ の行列が与えられる。これらの行列のある 3 つの行が good であるとは、「どの列についても、それら 3 行の値がすべて等しいか、すべて互いに異なる」という条件を満たすことを言う。

この行列の行を 5 つ選ぶ方法であって、その 5 行の中から good な 3 行の組を 2 つ以上選べるようなものの個数を求めよ。

制約

$1 \le N \le 1000$
$1 \le K \le 20$
どの 2 つの行も互いに異なる

考えたこと

0, 1, 2 のみかななる 3 値 $x, y, z$ が、すべて等しいかすべて互いに異なるという条件は、

$x + y + z \equiv 0 \pmod 3$

と言い換えられる。これはド典型らしい。このことから言えるのは、2 行を決めると、good になるべき残りの行の要素はただ一つに決まるということだ。

このことと「入力行列のどの 2 行も互いに異なる」という条件から、条件を満たす 5 行において、good な 3 行の組が 2 つとれるとき、

2 組の間で 1 つの行を共有している (行ベクトルを $x$ とする)
残りの 4 行は、2 ペアに分けられて
- 行ベクトル $a$ と $a, x$ から作れる行ベクトル
- 行ベクトル $b$ と $b, x$ から作れる行ベクトル

という形になることがわかる。 $a \neq b$ であるとき、これら 5 行は自動的にすべて互いに異なることになるので、ダブルカウントを気にする必要はない。

よって、次の方法によって $O(K N^{2} \log N)$ の計算量で解ける。

行 $i$ を固定して、これが good な 2 組が共有する行とする
残りの各行について
- 行 $i, j$ と good にある行ベクトルを計算し、それが行列に含まれるかを check する
- 含まれる個数を $n$ としたとき、 ${}_{\frac{n}{2}}\mathrm{C}_{2}$ 通りを加算していく

コード

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

int main() {
    int N, K;
    cin >> N >> K;
    vector<vector<int>> A(N, vector<int>(K));
    for (int i = 0; i < N; ++i) for (int j = 0; j < K; ++j) cin >> A[i][j];
    
    sort(A.begin(), A.end());
    
    long long res = 0;
    for (int i = 0; i < N; ++i) {
        long long num = 0;
        for (int j = 0; j < N; ++j) {
            if (j == i) continue;
            vector<int> sum(K, 0);
            for (int k = 0; k < K; ++k) {
                sum[k] = (3 - (A[i][k] + A[j][k]) % 3) % 3;
            }
            int it = lower_bound(A.begin(), A.end(), sum) - A.begin();
            if (A[it] == sum) ++num;
        }
        num /= 2;
        res += num * (num - 1) / 2;
    }
    cout << res << endl;
}

けんちょんの競プロ精進記録

競プロの精進記録や小ネタを書いていきます

Codeforces Round 824 (Div. 2) D. Meta-set (R1700)

問題概要

制約

考えたこと

コード