SRM410 Div2 Medium

500

グリッドと連結でない頂点は、最も大きい連結成分との間に辺を作る。

class AddElectricalWires {
public:
  int maxNewWires(vector <string> G, vector <int> C)
  {
    const int N = G.size();
    bool g[N][N];
    bool h[N][N];

    for (int i = 0; i < (int)N; ++i) {
      g[i][i] = h[i][i] = true;
      for (int j = 0; j < (int)N; ++j) {
        g[i][j] = (G[j][i] == '1');
        h[i][j] = g[i][j];
      }
    }

    for (int k = 0; k < (int)N; ++k) {
      for (int i = 0; i < (int)N; ++i) {
        for (int j = 0; j < (int)N; ++j) {
          g[i][j] = g[i][j] || (g[i][k] && g[k][j]);
        }
      }
    }

    int color[N];
    fill(color, color + N, -1);
    for (int i = 0; i < (int)C.size(); ++i) {
      int n = C[i];
      color[n] = n;
      for (int j = 0; j < (int)N; ++j) {
        int m = j;
        if (g[n][m]) color[m] = color[n];
      }
    }

    int mx = 0;
    int c = color[C.front()];
    for (int i = 0; i < (int)C.size(); ++i) {
      int tmp = count(color, color + N, C[i]);
      if (tmp > mx) {
        c = C[i];
        mx = tmp;
      }
    }

    int ret = 0;

    for (int i = 0; i < (int)N; ++i) {
      if (color[i] == -1) {
        ++ret;
        color[i] = c;
        h[i][c] = h[c][i] = true;
      }
    }

    for (int i = 0; i < (int)N; ++i) {
      for (int j = i + 1; j < (int)N; ++j) {
        if (!h[i][j] && color[i] == color[j]) ++ret;
      }
    }

    return ret;
  }
};