java-competitive/archive/2018.09/2018.09.30 - unsorted/TaskB.java

package chelper;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.HashSet;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.Queue;
import java.util.Set;
import java.util.TreeMap;

import io.InputReader;
import io.OutputWriter;
import misc.Algo;
import misc.Couple;
import misc.GCJSolution;


public class TaskB extends GCJSolution {

  public void solve(int testNumber, InputReader in, OutputWriter out) {
    wrapSolve(testNumber, in, out);
  }

  @Override
  public void solve(int testNumber) {


    int n = in.nextInt();
    int m = in.nextInt();

    Map<Integer, List<Couple<Integer>>> edges = new HashMap<>();

    for (int i = 0; i < n; i++) {
      edges.put(i, new ArrayList<>());
    }

    int special1 = -1;
    int special2 = -1;

    for (int i = 0; i < m; i++) {
      int a = in.nextInt() - 1;
      int b = in.nextInt() - 1;

      int c = in.nextInt();

      if (c == 0) {
        special1 = a;
        special2 = b;
      }

      edges.get(a).add(new Couple<>(b, c));
      edges.get(b).add(new Couple<>(a, c));
    }

//    long ans1 = solveBrute(n, m, edges);
//    out.println(ans1);
    long ans2 = solveSmart2(n, m, edges, special1, special2);
    out.println(ans2);

//    if (ans1 != ans2) {
////      out.println("case " + testNumber + " : " + n + " " + m);
//    }
  }

  long solveBrute(int n, int m, Map<Integer, List<Couple<Integer>>> edges) {
    int[][] distances = new int[n][n];

    for (int i = 0; i < n; i++) {
      Arrays.fill(distances[i], 100000000);
      distances[i][i] = 0;
    }

    for (int v = 0; v < n; v++) {
      for (Couple<Integer> edge : edges.get(v)) {
        int vv = edge.first;
        int c = edge.second;

        distances[v][vv] = Math.min(distances[v][vv], c);
        distances[vv][v] = Math.min(distances[vv][v], c);
      }
    }

    for (int i = 0; i < n; i++) {
      for (int j = 0; j < n; j++) {
        for (int k = 0; k < n; k++) {
          distances[i][j] = Math.min(distances[i][j], distances[i][k] + distances[k][j]);
        }
      }
    }

    int res = 0;
    int bestSum = Integer.MAX_VALUE;

    for (int mask = 1; mask < (1 << n) - 1; mask++) {
      boolean[] colors = new boolean[n];

      for (int i = 0; i < n; i++) {
        colors[i] = (mask & (1 << i)) > 0;
      }

      int sum = 0;

      for (int i = 0; i < n; i++) {
        int min = Integer.MAX_VALUE;

        for (int j = 0; j < n; j++) {
          if (colors[i] != colors[j]) {
            min = Math.min(min, distances[i][j]);
          }
        }

        sum += min;
      }

      if (sum < bestSum) {
        debug.println("new best sum " + sum);
        bestSum = sum;
        res = 0;
      }

      if (sum == bestSum) {
        res++;
      }

      if (sum == 2) {
        for (boolean color : colors) {
          debug.print(color ? '1' : '0');
        }
        debug.println();
      }
    }

    return res;
  }


  long solveSmart2(int n, int m, Map<Integer, List<Couple<Integer>>> edges, int special1, int special2) {
    TreeMap<Integer, Couple<Integer>> edgesSorted = new TreeMap<>();

    for (int v = 0; v < n; v++) {
      for (Couple<Integer> edge : edges.get(v)) {
        int vv = edge.first;
        int c = edge.second;

        edgesSorted.put(c, new Couple<>(v, vv));
      }
    }

    int components = 0;
    int zeroAdjacent = 0;

    boolean[] visited = new boolean[n];

    for (Map.Entry<Integer, Couple<Integer>> edge : edgesSorted.entrySet()) {
      int c = edge.getKey();

      int v = edge.getValue().first;
      int vv = edge.getValue().second;

      if (!visited[v] && !visited[vv]) {
        components++;
      }

      if (visited[v] != visited[vv]) {
        if (v == special1 || v == special2 || vv == special1 || vv == special2) {
          zeroAdjacent++;
        }
      }

      visited[v] = true;
      visited[vv] = true;
    }

    return 1L << (components + zeroAdjacent);
  }

  long solveSmart(int n, int m, Map<Integer, List<Couple<Integer>>> edges, int special1, int special2) {


    int components = 0;
    int specialAdd = 0;

    boolean[] visited = new boolean[n];

    if (special1 != -1) {
      List<Integer> specialStarting = new ArrayList<>();
      specialStarting.add(special1);
      specialStarting.add(special2);

      specialAdd = bfs(specialStarting, edges, visited) + 1;
    }

    for (int i = 0; i < n; i++) {
      if (!visited[i]) {
        List<Integer> starting = new ArrayList<>();
        starting.add(i);

        bfs(starting, edges, visited);

        components++;
      }
    }

    return 1L << (components + specialAdd);
  }

  int bfs(List<Integer> starting, Map<Integer, List<Couple<Integer>>> edges, boolean[] visited) {
    Queue<Integer> queue = new LinkedList<>(starting);

    Set<Integer> immediateNeighbours = new HashSet<>();

    for (int v : starting) {
      visited[v] = true;

      for (Couple<Integer> edge : edges.get(v)) {
        int vv = edge.first;
        int c = edge.second;

        immediateNeighbours.add(vv);
      }
    }
    immediateNeighbours.removeAll(starting);

    while (!queue.isEmpty()) {
      int v = queue.poll();

      for (Couple<Integer> edge : edges.get(v)) {
        int vv = edge.first;
        int c = edge.second;

        if (!visited[vv]) {
          visited[vv] = true;
          queue.add(vv);
        }

      }
    }

    return immediateNeighbours.size();
  }
}