Facebook Hacker Cup 2016 Round 2 - Problem D: Costly Labels

D.cc

#include <algorithm>
#include <iostream>
#include <utility>
#include <vector>
using namespace std;

#define FOR(i, a, b) for(int i = int(a); i < int(b); ++i)
#define REP(i, n)    FOR(i, 0, n)
#define SZ(c)        (int((c).size()))

/* Minimum-cost maximum flow implementation starts here. */
struct Edge { int src, dst, cap, cost, flow; };

const int inf = 999999999;
vector<Edge> edges;

/* Add an edge to the network with given maximum capacity and per-unit cost. */
void add_edge(int src, int dst, int cap, int cost) {
  edges.push_back({src, dst, cap, cost, 0});
}

/* Search for a minimum-cost augmenting path using Bellman-Ford.
   If an augmenting path is found, the flow is augmented and true is returned. */
bool augment(int V, int s, int t, int &cst, int &flw) {
  vector<int> cost(V, inf);
  vector<int> edge(V, 0);
  cost[s] = 0;
  bool changed;
  do {
    changed = false;
    REP(n, SZ(edges)) {
      const Edge &e = edges[n];
      if (e.flow < e.cap && cost[e.src] + e.cost < cost[e.dst]) {
        cost[e.dst] = cost[e.src] + e.cost;
        edge[e.dst] = n + 1;
        changed = true;
      }
      if (e.flow > 0 && cost[e.dst] - e.cost < cost[e.src]) {
        cost[e.src] = cost[e.dst] - e.cost;
        edge[e.src] = -(n + 1);
        changed = true;
      }
    }
  } while(changed);
  if (cost[t] == inf) return false;

  // Find min. flow over all edges
  int f = inf;
  for (int n = edge[t]; n != 0; ) {
    const Edge &e = edges[abs(n) - 1];
    f = min(f, n > 0 ? e.cap - e.flow : e.flow);
    n = edge[n > 0 ? e.src : e.dst];
  }

  // Augment flow
  for (int n = edge[t]; n != 0; ) {
    Edge &e = edges[abs(n) - 1];
    if (n > 0) e.flow += f, cst += f*e.cost;
    if (n < 0) e.flow -= f, cst -= f*e.cost;
    n = edge[n > 0 ? e.src : e.dst];
  }
  flw += f;
  return true;
}

/* Determine the minimum-cost maximum-flow (assuming all flows are initialized
   to zero first). This means flow is maximized first; of all maximal flows,
   a flow is selected with minimal cost.

   The return value is a pair of minimum cost and maximum flow.

   Complexity: O(max_flow*V*E) time, O(V) space
*/
pair<int,int> min_cost_max_flow(int s, int t) {
  if (s == t) return make_pair(0, inf);

  // Determine number of vertices
  int V = max(s, t);
  for (const Edge& edge : edges) V = max(V, max(edge.src, edge.dst));
  V += 1;

  int cst = 0, flw = 0;
  while (augment(V, s, t, cst, flw)) {};
  return make_pair(cst, flw);
}

// Real solution code starts here.
// These are the parameters to the problem.
static int N, K, P;
static int C[1000][1000];
static vector<vector<int>> adj;
static int memo[1000][31][31];

// Calculates the minimum cost of labelling the subtree rooted at vertex `i` if
// it gets color `col` and its parent has color `par_col`.
int Calc(int i, int col, int parent, int par_col) {
  // Check for a previously-computed result:
  int& m = memo[i][col + 1][par_col + 1];
  if (m >= 0) return m;

  // Calculate the possible costs of assigning each color to each child vertex.
  vector<vector<int>> costs;
  for (int j : adj[i]) if (j != parent) {
    vector<int> row(K);
    REP(k, K) row[k] = Calc(j, k, i, col) + C[j][k];
    costs.push_back(std::move(row));
  }

  // First option: pay the penalty P and then take the minimum cost for all subtrees.
  int res = P;
  for (vector<int>& row : costs)
    res += *min_element(row.begin(), row.end());

  // Second option: try to assign distinct colors to this node's parent/children.
  if (SZ(costs) + (parent >= 0) <= K) {
    edges.clear();
    REP(k, K) add_edge(0, 2 + k, 1, 0);
    REP(n, SZ(costs)) add_edge(2 + K + n, 1, 1, 0);
    REP(n, SZ(costs)) REP(k, K) if (k != par_col) add_edge(2 + k, 2 + K + n, 1, costs[n][k]);
    res = min(res, min_cost_max_flow(0, 1).first);
  }

  // Store and return the result.
  return m = res;
}

int Solve() {
  memset(memo, -1, sizeof(memo));
  int answer = 2000000000;
  REP(k, K) answer = min(answer, Calc(0, k, -1, -1) + C[0][k]);
  return answer;
}

int main() {
  int T = 0;
  cin >> T;
  REP(t, T) {
    cin >> N >> K >> P;
    REP(i, N) REP(j, K) cin >> C[i][j];
    adj.assign(N, {});
    REP(i, N - 1) {
      int a = 0, b = 0;
      cin >> a >> b;
      adj[a - 1].push_back(b - 1);
      adj[b - 1].push_back(a - 1);
    }
    cout << "Case #" << t + 1 << ": " << Solve() << endl;
  }
}