ABC016 C:友達の友達

https://atcoder.jp/contests/abc016/tasks/abc016_3

最初は隣接行列を使って解いた。ただし、節点同士が直接繋がっているかいないかという情報しかもたせていない。距離をもたせるなんて発想はなかった。思いついたとして、コードに落とし込めたかは怪しいところだ。

signed main() {
    int n, m; cin >> n >> m;
    int graph[10][10] = { 0 };
    for (int i = 0; i < m; i++) {
        int a, b; cin >> a >> b;
        a--, b--;
        graph[a][b] = 1; graph[b][a] = 1;
    }
 
    int used[10] = { 0 };
    for (int i = 0; i < n; i++) {
        memset(used, 0, sizeof(used));
        int ans = 0;
        for (int j = 0; j < n; j++) {
            if (i == j) continue;
            // iとjが友人なら、jの友人(iの友人の友人)を探索
            if (graph[i][j]) {
                for (int k = 0; k < n; k++) {
                    if (j == k || i == k) continue;
                    // iとkが直接の友人ではなく、kをまだ数えていない
                    if (graph[j][k]  && graph[i][k] == 0 && !used[k]) {
                        ans++;
                        used[k] = 1;
                    }
                }
            }
        }
        cout << ans << endl;
    }
}

うーむ、見辛い。バグりやすそうな実装だ。実際、テストケースがなかなか通らなくてデバッグに苦労した。
公式解説を見ると、ワーシャル・フロイド法を使えばいいとのこと。いい機会だ。今までのように概要だけサーっと流すのではなく、実際に自分が使うことを考えてコードに落としこむことにした。

まあ実のところはネット上のあらゆる場所からパクりながら書いたわけだが。

ワーシャル・フロイド法そのものに関する解説は、以下の記事が非常にわかりやすい。

http://doriven.hatenablog.com/entry/2013/12/23/044921

class Edge {
private:
    int m_from, m_to, m_cost;
 
public:
    Edge(int f, int t, int c) :
        m_from(f), m_to(t), m_cost(c) {}
 
    int from() const { return m_from; }

    int to() const { return m_to; }

    int cost() const { return m_cost; }
 
    Edge& operator=(const int x) {
        m_cost = x;
        return *this;
    }
};
 
using Edges = vector<Edge>;
using Graph = vector<Edges>;
 
// 全点間最短路を求める。到達できない要素にはinfが格納されている前提。
void warshall_floyd(Graph& graph, const int inf) {
    const int v = graph.size();
    for (int k = 0; k < v; k++) {
        for (int i = 0; i < v; i++) {
            for (int j = 0; j < v; j++) {
                //if (i == j || i == k || j == k) continue; // 自分と同じノードに対する距離に0を格納していたらこの処理は不要。
                if (graph[i][k].cost() == inf || graph[k][j].cost() == inf) continue;
                graph[i][j] = min(graph[i][j].cost(), graph[i][k].cost() + graph[k][j].cost());
            }
        }
    }
}
 
// 隣接行列を生成
Graph gen_graph(const int n, const int inf) {
    Graph g(n);
    for (int i = 0; i < n; i++) {
        for (int j = 0; j < n; j++) {
            // 最初はすべてのノードが繋がっていないとする。
            // 自分自身に対する距離は0、繋がっていないノードへの距離はinf
            int cost = (i == j) ? 0 : inf;
            g[i].push_back(Edge(i, j, cost));
        }
    }
    return g;
}

signed main() {
    int n, m; cin >> n >> m;
    auto graph = gen_graph(n, INF);
    for (int i = 0; i < m; i++) {
        int a, b; cin >> a >> b; a--, b--;
        graph[a][b] = graph[b][a] = 1;
    }
    warshall_floyd(graph, INF);
 
    for (int i = 0; i < n; i++) {
        int ans = 0;
        for (int j = 0; j < n; j++) {
            // 友人の友人 -> 距離が2
            if (graph[i][j].cost() == 2) ans++;
        }
        cout << ans << endl;
    }   
}

どのコードがどういう働きをしているかについて理解しやすくなった。
早解きがレート更新で重要な競プロにおいて、ここまでの実装はいささか大げさのような気がしなくもないが、おれのような雑魚はまず問題を解けるかどうかの方が重要だ。コードは見やすければ見やすいほど良い。

ついでに、いつだかのコンテストで「B問題でワーシャル・フロイドが出やがった」と一部のtwitter界隈で話題になったアレを

https://guri6cerin.hatenablog.com/entry/2019/02/10/AtCoder_%E3%80%8C%E3%81%BF%E3%82%93%E3%81%AA%E3%81%AE%E3%83%97%E3%83%AD%E3%82%B3%E3%83%B3_2019%E3%80%8D

ワーシャル・フロイドで解いて見た。

signed main() {
    auto graph = gen_graph(4, INF);
    for (int i = 0; i < 3; i++) {
        int a, b; cin >> a >> b; a--, b--;
        graph[a][b] = graph[b][a] = 1;
    }
    warshall_floyd(graph, INF);
 
    bool ok = false;
    for (int i = 0; i < 4; i++) {
        for (int j = 0; j < 4; j++) {
            if (graph[i][j].cost() == 3) {
                ok = true;
            }
        }
    }
    cout << (ok ? "YES" : "NO") << endl;
}

・・・うん、まあこれこそオーバーキルというものだろう。早すぎた最適化というか。ちょっと違うか。