atcoderのabc254_eでTLE

c++ 競技プログラミング

こちらの問題で実行時間制限を超えてしまったのだが、普段あまり意識していなかった部分で引っかかっていたのでメモ。
https://atcoder.jp/contests/abc254/tasks/abc254_e

問題の概要は以下

  • グラフが与えられる(頂点数N=1.5e5)
  • クエリが与えられる(クエリ数Q=1.5e5)
  • クエリで与えられた頂点の近傍kまでの頂点番号の和を取る

気をつける点としては計算量で

  • 対象の頂点の近傍kまでの頂点のみ探索する必要がある
    • 頂点の次数は3以下というのが問題で与えられているため、途中で探索を打ち切れば一回あたりの探索頂点数は100以下で済む
  • なので全体の計算量は 100Nには収まり無事制限時間内に終わる

という形。なので素直に幅優先探索を実装して条件に合わせて探索を打ち切るようにしたのだがそれでもTLEとなってしまった。

コードはこんな感じ。

using Graph = vector<vector<int>>;
using P = pair<int,int>;

ll bfs(const Graph &G, int x, int k) {
    int N = G.size();
    deque<P> que;
    que.push_back({0, x});

    vector<int> dist(N, -1);
    dist[x] = 0;

    ll ret = x+1;
    int count = 0;
    while (!que.empty()) {
        auto &p = que.front();
        que.pop_front();

        int d = p.first;
        int v = p.second;

        // 元の頂点からの距離がk以上になった場合には探索終了
        if (d >= k) break;

        for (auto nv : G[v]) {
            ++count;
            if (dist[nv] != -1) continue;
            dist[nv] = d+1;
            ret += nv+1;
            que.emplace_back(d+1, nv);
        }
    }

    return ret;
}

void _main() {
    int N, M;
    cin >> N >> M;

    Graph G(N);
    REP(i, M) {
        int a, b; cin >> a >> b;
        --a; --b;
        G[a].push_back(b);
        G[b].push_back(a);
    }

    int Q;
    cin >> Q;
    REP(i, Q) {
        int x, k; cin >> x >> k;
        --x;
        cout << bfs(G, x, k) << endl;
    }

}

int main() {
    _main();
    return 0;
}

結局コンテスト終了までにバグは見つけられなかったので、終了後に解説のコードと比較しつつ考えたところ、やはり探索部分についての考え方や実装は問題なさそうだった。
https://atcoder.jp/contests/abc254/editorial/4052

その後も気づくまでに時間がかかったのだが、最終的にどこが問題だったかというとbfsの関数内でdistのvectorを初期化していることだった。

vector<int> dist(N, -1);

ここが普段私が意識していなかった部分だったのだが考えてみれば当然で、メモリ確保とその初期値設定には大きさに応じた処理数が必要になる。今回の問題の場合これを行う関数がクエリ数分呼ばれる形となっていたため、合計で計算量がNQ > 1e10となっていた。

自分にとっての盲点に気づけて大変勉強になりました。