Line data    Source code

       1              : // ***************************************************************
       2              : // SPDX-FileCopyrightText: Copyright 2024 Ricardo Montañana Gómez
       3              : // SPDX-FileType: SOURCE
       4              : // SPDX-License-Identifier: MIT
       5              : // ***************************************************************
       6              : 
       7              : #include <sstream>
       8              : #include <vector>
       9              : #include <list>
      10              : #include "Mst.h"
      11              : /*
      12              :     Based on the code from https://www.softwaretestinghelp.com/minimum-spanning-tree-tutorial/
      13              : 
      14              : */
      15              : 
      16              : namespace bayesnet {
      17          116 :     Graph::Graph(int V) : V(V), parent(std::vector<int>(V))
      18              :     {
      19          434 :         for (int i = 0; i < V; i++)
      20          376 :             parent[i] = i;
      21           58 :         G.clear();
      22           58 :         T.clear();
      23           58 :     }
      24         1152 :     void Graph::addEdge(int u, int v, float wt)
      25              :     {
      26         1152 :         G.push_back({ wt, { u, v } });
      27         1152 :     }
      28         5326 :     int Graph::find_set(int i)
      29              :     {
      30              :         // If i is the parent of itself
      31         5326 :         if (i == parent[i])
      32         2304 :             return i;
      33              :         else
      34              :             //else recursively find the parent of i
      35         3022 :             return find_set(parent[i]);
      36              :     }
      37          318 :     void Graph::union_set(int u, int v)
      38              :     {
      39          318 :         parent[u] = parent[v];
      40          318 :     }
      41           58 :     void Graph::kruskal_algorithm()
      42              :     {
      43              :         // sort the edges ordered on decreasing weight
      44         4484 :         stable_sort(G.begin(), G.end(), [](const auto& left, const auto& right) {return left.first > right.first;});
      45         1210 :         for (int i = 0; i < G.size(); i++) {
      46              :             int uSt, vEd;
      47         1152 :             uSt = find_set(G[i].second.first);
      48         1152 :             vEd = find_set(G[i].second.second);
      49         1152 :             if (uSt != vEd) {
      50          318 :                 T.push_back(G[i]); // add to mst std::vector
      51          318 :                 union_set(uSt, vEd);
      52              :             }
      53              :         }
      54           58 :     }
      55              : 
      56          318 :     void insertElement(std::list<int>& variables, int variable)
      57              :     {
      58          318 :         if (std::find(variables.begin(), variables.end(), variable) == variables.end()) {
      59          318 :             variables.push_front(variable);
      60              :         }
      61          318 :     }
      62              : 
      63           58 :     std::vector<std::pair<int, int>> reorder(std::vector<std::pair<float, std::pair<int, int>>> T, int root_original)
      64              :     {
      65              :         // Create the edges of a DAG from the MST
      66              :         // replacing unordered_set with list because unordered_set cannot guarantee the order of the elements inserted
      67           58 :         auto result = std::vector<std::pair<int, int>>();
      68           58 :         auto visited = std::vector<int>();
      69           58 :         auto nextVariables = std::list<int>();
      70           58 :         nextVariables.push_front(root_original);
      71          434 :         while (nextVariables.size() > 0) {
      72          376 :             int root = nextVariables.front();
      73          376 :             nextVariables.pop_front();
      74         1328 :             for (int i = 0; i < T.size(); ++i) {
      75          952 :                 auto [weight, edge] = T[i];
      76          952 :                 auto [from, to] = edge;
      77          952 :                 if (from == root || to == root) {
      78          318 :                     visited.insert(visited.begin(), i);
      79          318 :                     if (from == root) {
      80          212 :                         result.push_back({ from, to });
      81          212 :                         insertElement(nextVariables, to);
      82              :                     } else {
      83          106 :                         result.push_back({ to, from });
      84          106 :                         insertElement(nextVariables, from);
      85              :                     }
      86              :                 }
      87              :             }
      88              :             // Remove visited
      89          694 :             for (int i = 0; i < visited.size(); ++i) {
      90          318 :                 T.erase(T.begin() + visited[i]);
      91              :             }
      92          376 :             visited.clear();
      93              :         }
      94           58 :         if (T.size() > 0) {
      95            0 :             for (int i = 0; i < T.size(); ++i) {
      96            0 :                 auto [weight, edge] = T[i];
      97            0 :                 auto [from, to] = edge;
      98            0 :                 result.push_back({ from, to });
      99              :             }
     100              :         }
     101          116 :         return result;
     102           58 :     }
     103              : 
     104           58 :     MST::MST(const std::vector<std::string>& features, const torch::Tensor& weights, const int root) : features(features), weights(weights), root(root) {}
     105           58 :     std::vector<std::pair<int, int>> MST::maximumSpanningTree()
     106              :     {
     107           58 :         auto num_features = features.size();
     108           58 :         Graph g(num_features);
     109              :         // Make a complete graph
     110          376 :         for (int i = 0; i < num_features - 1; ++i) {
     111         1470 :             for (int j = i + 1; j < num_features; ++j) {
     112         1152 :                 g.addEdge(i, j, weights[i][j].item<float>());
     113              :             }
     114              :         }
     115           58 :         g.kruskal_algorithm();
     116           58 :         auto mst = g.get_mst();
     117          116 :         return reorder(mst, root);
     118           58 :     }
     119              : 
     120              : }