TopCoder SRM 695 Div1 Medium: BearKRoads

解法

happy が最大となる辺の、大きい方の端は必ず選択される。

コード

import java.util.ArrayList;

public class BearKRoads {
  private int N;
  ArrayList<Integer>[] graph;
  int[] a, b;

  int dfs(int[] x, int k) {
    if (k == 0) return 0;
    ArrayList<int[]> ab = new ArrayList<>();
    int max = 0;
    for (int i = 0; i < a.length; i++) {
      int h = x[a[i]] + x[b[i]];
      if (h == 0) continue;
      if (max > h) continue;
      if (max < h) {
        ab.clear();
        max = h;
      }
      ab.add(new int[]{a[i], b[i]});
    }

    int ret = 0;
    for (int[] e : ab) {
      int p = e[0];
      if (x[e[0]] < x[e[1]]) p = e[1];
      for (int u : graph[p]) {
        int t1 = x[p];
        int t2 = x[u];
        int t = t1 + t2;
        x[p] = 0;
        x[u] = 0;
        t += dfs(x, k - 1);
        x[p] = t1;
        x[u] = t2;
        ret = Math.max(ret, t);
      }
    }
    return ret;
  }

  public int maxHappy(int[] x, int[] a, int[] b, int K) {
    N = x.length;
    this.a = a;
    this.b = b;
    graph = new ArrayList[N];
    for (int i = 0; i < N; i++) graph[i] = new ArrayList<>();
    for (int i = 0; i < a.length; i++) {
      graph[a[i]].add(b[i]);
      graph[b[i]].add(a[i]);
    }
    return dfs(x, K);
  }

}