AOJ 2309 Vector Compression

問題文

まず問題読解が難しかった．
http://judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=2309

問題概要

N次元のベクトルが M 個与えられる．これを v[1], ..., v[M] とする．
M 個のベクトルを好きな順番に並べ替え，その順に記録していく．記録する際にコストがかかる．コストは，|v[i]|^2 か，それ以前に記録していたベクトルの中から好きなもの v[j] を選んで，|v[i] - r * v[j]|^2 のどちらかを選択できる．ただし，r は記録のたびに好きな値を適当にとって良いとする．
この時，すべてのベクトルを記録するのにかかるコストを最小化せよ．

・制約
1 <= N, M <= 100

解法

v[i] を記録するときに v[j] を使う，というのを有向辺 (j, i) （向きに注意）であらわし，そのコストを |v[i] - r * v[j]|^2 の最小値とする．
また，v[i] を直接使う場合は，有向辺 (M, i) で表し，そのコストを |v[i]|^2 とする．
すると，この問題はこうして出来たグラフの最小有向全域木を求めよ，と言い換えることができる．
r の求め方だが，幸い |v[i] - r * v[j]|^2 は明らかに下に凸であり，かつ r で微分するのも簡単なので，極値が簡単に求められる（平方完成してもいい）．

計算量は O(nm^2)．

ソースコード

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

using weight = double;
constexpr weight INF = 1e9;

constexpr weight eps = 1e-8;

struct edge {
    int from, to;
    weight cost;
    
    bool operator<(edge const& e) const {
        return cost < e.cost;
    }
};

using edges = std::vector<edge>;
using graph = std::vector<edges>;

void add_edge(graph& g, int from, int to, weight cost) {
    g[from].push_back(edge{from, to, cost});
}

int scc(std::vector<std::vector<int>>& G, std::vector<int>& cmp) {
    int V = G.size();
    std::vector<std::vector<int>> g(V), rg(V);
    std::vector<bool> used(V, false);
    std::vector<int> vs;
    cmp.resize(V);
    for(int i = 0; i < V; ++i) {
        for(auto to : G[i]) {
            g[i].push_back(to);
            rg[to].push_back(i);
        }
    }
    std::function<void(int)> dfs = [&g, &vs, &used, &dfs](int v) {
        used[v] = true;
        for(auto i : g[v]) {
            if(!used[i]) {
                dfs(i);
            }
        }
        vs.push_back(v);
    };
    std::function<void(int, int)> rdfs = [&rg, &cmp, &used, &rdfs](int v, int k) {
        used[v] = true;
        cmp[v] = k;
        for(int i : rg[v]) {
            if(!used[i]) {
                rdfs(i, k);
            }
        }
    };
    for(int v=0; v<V; ++v) {
        if(!used[v]) {
            dfs(v);
        }
    }
    std::fill(used.begin(), used.end(), false);
    int k = 0;
    for(int i=vs.size()-1; i>=0; --i) {
        if(!used[vs[i]]) {
            rdfs(vs[i], k++);
        }
    }
    return k;
}

std::vector<std::vector<int>> build_graph(std::vector<std::vector<int>> const& g, std::vector<int> const& cmp, int K) {
    int V = g.size();
    std::vector<std::set<int>> s(K);
    std::vector<std::vector<int>> res(K);
    for(int i = 0; i < V; ++i) {
        for(auto to : g[i]) {
            s[cmp[i]].insert(cmp[to]);
        }
    }
    for(int i = 0; i < K; ++i) {
        for(auto j : s[i]) {
            if(i != j) {
                res[i].push_back(j);
            }
        }
    }
    return res;
}


weight minimum_spanning_rooted_arborescence(graph const& g, int root, weight sum = 0) {
    int const n = g.size();

    std::vector<int> rev(n, -1);
    std::vector<weight> cost(n, INF);
    for(int i = 0; i < n; ++i) {
        for(auto& e : g[i]) {
            if(e.cost < cost[e.to]) {
                cost[e.to] = e.cost;
                rev[e.to] = i;
            }
        }
    }
    for(int i = 0; i < n; ++i) {
        if(i != root && rev[i] == -1) { // not exists
            return INF;
        }
    }

    std::vector<std::vector<int>> g2(n);
    for(int i = 0; i < n; ++i) {
        if(root == i) {
            continue;
        }
        g2[rev[i]].push_back(i);
        sum += cost[i];
    }
    std::vector<int> cmp(n);
    int const K = scc(g2, cmp);
    auto nxt = build_graph(g2, cmp, K);
    if(nxt.size() == n) {
        return sum;
    }

    graph ng(nxt.size());
    for(int i = 0; i < n; ++i) {
        for(auto& e : g[i]) {
            if(cmp[i] == cmp[e.to]) {
                continue;
            }
            ng[cmp[i]].push_back(edge{cmp[i], cmp[e.to], e.cost - cost[e.to]});
        }
    }

    return minimum_spanning_rooted_arborescence(ng, cmp[root], sum);
}

int main() {
    int N, M;
    cin >> N >> M;
    vector<vector<double>> v;
    for(int i = 0; i < M; ++i) {
        double t = 0;
        vector<weight> w(N);
        for(int j = 0; j < N; ++j) {
            cin >> w[j];
            t += w[j] * w[j];
        }
        if(t > eps) {
            v.push_back(w);
        }
    }
    M = v.size();
    graph g(M + 1);
    int const root = M;
    for(int i = 0; i < M; ++i) {
        for(int j = 0; j < M; ++j) {
            if(i == j) {
                continue;
            }
            double prod = 0, dom = 0;
            for(int k = 0; k < N; ++k) {
                prod += v[i][k] * v[j][k];
                dom += v[j][k] * v[j][k];
            }
            double r = prod / dom;
            double cost = 0;
            for(int k = 0; k < N; ++k) {
                cost += (v[i][k] - r * v[j][k]) * (v[i][k] - r * v[j][k]);
            }
            add_edge(g, j, i, cost);
        }
        double d = 0;
        for(int j = 0; j < N; ++j) {
            d += v[i][j] * v[i][j];
        }
        add_edge(g, root, i, d);
    }
    cout << fixed << setprecision(10) << minimum_spanning_rooted_arborescence(g, root) << endl;
}

感想

最初フローかなあと思っていて全然解けず困っていた．
最小有向全域木を書いたのははじめてなのでいい勉強になりました．