content: package by.andd3dfx.common;
  
  import java.util.Arrays;
  import java.util.HashSet;
  import java.util.Set;
  import java.util.stream.Collectors;
  
  /**
   * <pre>
   * В королевстве есть несколько городов, соединенных дорогами. Города называются соседними, если между ними есть прямая
   * дорога. Первоначально чума вспыхнула лишь в паре городов, однако каждую неделю соседние города чумного города также
   * подвергались заражению и становились чумными.
   * Задача определить через сколько недель все королевство будет заражено.
   *
   * Input:
   * - cities_amount: n
   * - roads: [[n_i, n_j], ...]
   * - infected: [m_1, m_2, ...]
   *
   * Output:
   * weeks_count
   * </pre>
   *
   * @see <a href="https://youtu.be/Ei1uCCD_Iqg">Video solution (initial)</a>
   * @see <a href="https://youtu.be/d9v9DS1YVtk">Video solution (final)</a>
   */
  public class InfectionDistribution {
  
      public static int weeksToInfectAllCities(int citiesAmount, int[][] roads, int[] infected) {
          var adjMatrix = new boolean[citiesAmount][citiesAmount];    // adjacency matrix
          Arrays.stream(roads).forEach(road -> {
              adjMatrix[road[0]][road[1]] = true;
              adjMatrix[road[1]][road[0]] = true;
          });
  
          var weeksAmount = 0;
          var infectedCities = Arrays.stream(infected).boxed().collect(Collectors.toSet());
  
          while (infectedCities.size() < citiesAmount) {
              var newCities = determineNewInfectedCities(infectedCities, adjMatrix);
              if (newCities.isEmpty()) {
                  return -1;
              }
  
              infectedCities.addAll(newCities);
  
              weeksAmount++;
          }
  
          return weeksAmount;
      }
  
      private static Set<Integer> determineNewInfectedCities(Set<Integer> infectedCities, boolean[][] adjMatrix) {
          var n = adjMatrix.length;
          var newCities = new HashSet<Integer>();
          for (var city : infectedCities) {
              for (int i = 0; i < n; i++) {
                  if (adjMatrix[city][i] && !infectedCities.contains(i)) {
                      newCities.add(i);
                  }
              }
          }
  
          return newCities;
      }
  }
  