topsort c++

topsort.cpp

struct TopSortResult {
  bool ok;

  // If `ok`, contains the topologically sorted node indices.
  // Otherwise, contains indices of a detected cycle.
  std::vector<std::size_t> nodes;
};

// Topologically sorts dependencies.
TopSortResult topsort(const std::vector<std::vector<int>> &edges) {
  TopSortResult result;
  enum State { UNVISITED = 0, VISITING = 1, VISITED = 2 };
  const std::size_t num_vertices = edges.size();
  std::vector<State> state(num_vertices, UNVISITED);
  std::vector<int> stack;
  std::vector<int> prev(num_vertices, -1);  // for cycle detection
  stack.reserve(num_vertices);
  for (int start = 0; start < num_vertices; ++start) {
    if (state[start] == VISITED) continue;
    stack.push_back(start);
    while (!stack.empty()) {
      const int v = stack.back();

      const int pos_v = v < 0 ? -v - 1 : v;
      if (v < 0) {
        // Finished visiting:
        stack.pop_back();
        state[pos_v] = VISITED;
        result.nodes.push_back(pos_v);
        continue;
      }
      if (state[v] == VISITED) {
        stack.pop_back();
        continue;
      }

      if (state[v] == VISITING) {
        // Cycle detected:
        result.ok = false;
        result.nodes.clear();
        result.nodes.push_back(v);
        int cur = v;
        do {
          cur = prev[cur];
          if (cur == -1) break;
          result.nodes.push_back(cur);
        } while (cur != v);
        return result;
      }

      // When `-v - 1` is popped, we've finished visiting `v`.
      stack.back() = -v - 1;
      state[v] = VISITING;
      for (int u : edges[v]) {
        if (state[u] != VISITED) {
          prev[u] = v;
          stack.push_back(u);
        }
      }
    }
  }
  result.ok = true;
  return result;
}