Main.java

import java.util.*;

class Main {
  public static void main(String[] args) {
    case0();
    case1();
  }

  public static void case1() {
    int k = 2;
    int[] time = { 2, 3, 4, 1 };
    Map<Integer, Set<Integer>> edges = new HashMap<>();
    // 0 -> 1 -> 3
    //      2
    for(int i = 0; i < 4; i++) {
      Set<Integer> deps = new HashSet<>();
      if(i == 0) {
        deps.add(1);
      } else if(i == 1) {
        deps.add(3);
      }
      edges.put(i, deps);
    }
    Main solver = new Main();
    System.out.println(solver.solve(edges, 2, time)); // 6
  }


  public static void case0() {
    int k = 2;
    int[] time = { 2, 3, 4, 1 };
    Map<Integer, Set<Integer>> edges = new HashMap<>();
    // 0 -> 1 -> 3
    //      |
    //      +--> 2
    for(int i = 0; i < 4; i++) {
      Set<Integer> deps = new HashSet<>();
      if(i == 0) {
        deps.add(1);
      } else if(i == 1) {
        deps.add(3);
        deps.add(2);
      }
      edges.put(i ,deps);
    }
    Main solver = new Main();
    System.out.println(solver.solve(edges, 2, time)); // 9
  }

  /**
  return: minimum time that can complete all the tasks
  */
  int solve(
    Map<Integer, Set<Integer>> edges, // tasks dependency graph, x -> y, means x depends on y, y has to complete before x start
    int k,     // number of processors that can process tasks
    int[] time // time needed for each task on one processor
  ) {
    PriorityQueue<int[]> processor = new PriorityQueue<>((a, b) -> {
      return a[1] - b[1];
    }); // task_id, completion_time
    Set<Integer> completableTasks = new HashSet<>();
    Set<Integer> completedTasks = new HashSet<>();
    Map<Integer, Set<Integer>> backwardEdges = new HashMap<>();
    int[] taskOutDegreeCount = new int[edges.size()];
    for(Map.Entry<Integer, Set<Integer>> entry: edges.entrySet()) {
      if(entry.getValue().isEmpty()) completableTasks.add(entry.getKey());
      int from = entry.getKey();
      taskOutDegreeCount[from] = edges.get(from).size();
      for(int to: entry.getValue()) {
        if(!backwardEdges.containsKey(to)) backwardEdges.put(to, new HashSet<>());
        backwardEdges.get(to).add(from);
      }
    }

    int maxTime = 0;

    while(completedTasks.size() != edges.size()) {
      int timeElapsed = 0;
      while(
        (processor.size() == k || completableTasks.isEmpty())
        && !processor.isEmpty()
        ) {
        int[] completed = processor.remove();
        System.out.println(String.format("completed task %d at %d", completed[0], completed[1]));
        maxTime = Math.max(maxTime, completed[1]);
        completedTasks.add(completed[0]);
        completableTasks.remove(completed[0]);
        timeElapsed = completed[1];

        Set<Integer> parents = backwardEdges.get(completed[0]);
        if(parents != null) {
          for(int p: parents) {
            taskOutDegreeCount[p]--;
            if(taskOutDegreeCount[p] == 0) completableTasks.add(p);
          }
        }
      }
      
      if(completableTasks.isEmpty()) break;

      int candidate = completableTasks.iterator().next();
      int[] task = { candidate, timeElapsed + time[candidate] };
      System.out.println("added task " + task[0]);
      processor.add(task);
      completableTasks.remove(candidate);
    }

    while(!processor.isEmpty()) {
      int[] topItem = processor.remove();
      maxTime = Math.max(maxTime, topItem[1]);
    }

    return maxTime;

  }
}