Instantly share code, notes, and snippets.

Embed
What would you like to do?
Facebook Hacker Cup 2016 Round 2 - Problem D: Costly Labels
#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;
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment