Installation des réservoirs et robinets dans un quartier - Programmation compétitive

#include <iostream>
#include <cstring>
#include <vector>
#include <stack>
#include <set>
#include <algorithm>
#include <climits>

using namespace std;

// Classe pour représenter le graphe
class GraphAdjList {
    private:
        vector<vector<int>> voisins;
        vector<vector<int>> diametres;
        int nbMaisons;
        
    public:
        GraphAdjList(const vector<vector<int>> &tuyaux, int nb_maisons) {
            nbMaisons = nb_maisons;
            voisins.resize(nb_maisons);
            diametres.resize(nb_maisons);
            
            // tuyaux[i][0] => f_i ; tuyaux[i][1] => t_i ; tuyaux[i][2] => d_i
            for (int i = 0; i < tuyaux.size(); i++) {
                // ajouter t_i dans la liste d'adjacence de f_i
                voisins[tuyaux[i][0]].push_back(tuyaux[i][1]);
                // le diamètre entre f_i et t_i est d_i
                diametres[tuyaux[i][0]].push_back(tuyaux[i][2]);
            }
        }

        const vector<vector<int>> & getVoisins() { return voisins; }
        const vector<vector<int>> & getDiametres() { return diametres; }
        int getNbMaisons() { return nbMaisons; }
};

// Parcours en largeur pour trouver toutes les maisons connectées à partir de 'depart'
vector<int> parcours_largeur(GraphAdjList &reseau, int depart) {
    vector<int> connected;
    bool visited[reseau.getNbMaisons()];
    memset(visited, false, sizeof(visited));
    
    stack<int> pile;
    pile.push(depart);
    
    while (!pile.empty()) {
        int current = pile.top();
        pile.pop();
        
        if (!visited[current]) {
            connected.push_back(current);
            visited[current] = true;

            for (int v : reseau.getVoisins()[current]) {
                pile.push(v);
            }
        }
    }
    return connected;
}

void solution(vector<vector<int>> tuyaux, int n) {
    // Créer le graphe (les maisons sont numérotées de 1 à n, on utilise n+1 pour simplifier)
    GraphAdjList reseau(tuyaux, n + 1);

    // Stocker tous les noeuds qui ont des arêtes entrantes (robinets potentiels)
    set<int> robinets;
    for (auto &tuyau : tuyaux) {
        robinets.insert(tuyau[1]);
    }

    // Stocker les réservoirs (noeuds avec arêtes sortantes mais sans arêtes entrantes)
    set<int> reservoires;
    for (auto &tuyau : tuyaux) {
        if (robinets.count(tuyau[0]) == 0) {
            reservoires.insert(tuyau[0]);
        }
    }

    // Afficher le nombre de réservoirs
    cout << reservoires.size() << endl;

    // Pour chaque réservoir, trouver la chaîne connectée
    for (int depart : reservoires) {
        vector<int> connected = parcours_largeur(reseau, depart);
        
        // Trouver le diamètre minimum dans cette chaîne
        int mini = INT_MAX;
        for (int i = 0; i < connected.size() - 1; i++) {
            vector<int> w = reseau.getDiametres()[connected[i]];
            int min_diam = *min_element(w.begin(), w.end());
            mini = min(min_diam, mini);
        }
        
        // Afficher réservoir, robinet (dernier élément), diamètre minimum
        cout << connected[0] << " " << connected[connected.size() - 1] << " " << mini << endl;
    }
}

int main() {
    int n, p, val;
    cin >> n >> p;
    
    vector<vector<int>> tuyaux(p);
    for (int i = 0; i < p; i++) {
        for (int j = 0; j < 3; j++) {
            cin >> val;
            tuyaux[i].push_back(val);
        }
    }

    solution(tuyaux, n);
    
    return 0;
}