G - Triangle 解説 by kusuninu

bitsetを使用しない方法

\(A_{i,j}=A_{j,k}=A_{i,k}=1\)を満たす、\(1\le i < j < k \le N\)である整数の組\((i,j,k)\)を愚直に数えることを考えます。以下のようなコードで計算できます。

long long ans = 0;
for(int i = 0; i < N; i++) for(int j = i + 1; j < N; j++) for(int k = j + 1; k < N; k++) {
    if(A[i][j] && A[j][k] && A[i][k]) ans++;
}

しかし、このコードだと、論理積演算子&&の計算を最大で\(3000^3/6 \times 2=9\times 10^9 \)回行わなければならず、TLEしてしまいます。ところが、&&をビット毎AND演算子の&に書き換えた以下のコードは制限時間に間に合います。

long long ans = 0;
for(int i = 0; i < N; i++) for(int j = i + 1; j < N; j++) for(int k = j + 1; k < N; k++) {
    ans += A[i][j] & A[j][k] & A[i][k];
}

これは、条件分岐のオーバーヘッドとコンパイラによる自動ベクトル化(SIMD)が原因です。論理積演算子&&は短絡評価を行うため、条件分岐により処理に時間がかかります。短絡評価とは、X&& Yという式において、Xがfalseである場合、Yを評価せずに式をfalseとすることです。Xがtrueの場合とfalseの場合で処理が分岐するため、条件分岐によるオーバーヘッドが発生します。 一方、ビット毎AND演算子&は短絡評価を行わず、常に両方を計算します。そのため、条件分岐のコストが発生しません。条件分岐が発生しないため、コンパイラによってループ内の&の計算が複数のデータを一度に処理するSIMD命令を使って最適化されます。これにより、256bitの命令では8bitのデータを32個同時に処理することができ、処理速度が大幅に高速化してACすることができます。なお、この最適化はGCCよりもClangの方が積極的に行う傾向があります。

その他、以下のような実装の工夫によってさらに高速化することができます。

  • A[k][j]ではなくA[j][k]のようにアクセスすることで、シーケンシャルアクセスによりキャッシュに乗りやすくする

  • \(ijk\)ループ内で、\(A_{i,j}=0\)の場合を枝狩りする

  • \(A_{i,j}\)の型として、4Byteのintの代わりに1Byte のcharを使うことで、SIMDの並列数を増やす

  • 高速な入出力を利用する

  • メモリを静的に確保する

実装例(C++23 (Clang 21.1.0), 766ms) https://atcoder.jp/contests/abc258/submissions/77060408

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

constexpr int MAX_N = 3000;

char A[MAX_N][MAX_N];
char S[MAX_N][MAX_N];

int main() {
    int N;
    cin >> N;
    for(int i = 0; i < N; i++) {
        scanf("%s",S[i]);
    }
    for(int i = 0; i < N; i++) for(int j = 0; j < N; j++) A[i][j] = S[i][j] - '0';
    long long ans = 0;
    for(int i = 0; i < MAX_N;i++) for(int j = i + 1; j < MAX_N;j++){
        if(!A[i][j]) continue;
        for(int k = j + 1; k < MAX_N; k++){
            ans += A[j][k] & A[i][k];
        }
    }
    cout << ans << endl;
}

投稿日時:
最終更新: