tests/minimum_steiner_tree/restore.test.cpp

View this file on GitHub
Last update: 2025-05-10 19:26:48+09:00
Problem: https://judge.yosupo.jp/problem/minimum_steiner_tree

Depends on

Code

// competitive-verifier: PROBLEM https://judge.yosupo.jp/problem/minimum_steiner_tree

#include <iostream>
#include <vector>
#include "tools/join.hpp"
#include "tools/minimum_steiner_tree.hpp"

using ll = long long;

int main() {
  std::cin.tie(nullptr);
  std::ios_base::sync_with_stdio(false);

  int N, M;
  std::cin >> N >> M;
  tools::minimum_steiner_tree<ll> graph(N);
  for (int i = 0; i < M; ++i) {
    int U, V, W;
    std::cin >> U >> V >> W;
    graph.add_edge(U, V, W);
  }
  int K;
  std::cin >> K;
  std::vector<int> X(K);
  for (auto& X_i : X) std::cin >> X_i;

  const auto qr = graph.query<true>(X);
  std::cout << qr.cost << ' ' << qr.edge_ids.size() << '\n';
  std::cout << tools::join(qr.edge_ids, " ") << '\n';

  return 0;
}

#line 1 "tests/minimum_steiner_tree/restore.test.cpp"
// competitive-verifier: PROBLEM https://judge.yosupo.jp/problem/minimum_steiner_tree

#include <iostream>
#include <vector>
#line 1 "tools/join.hpp"



#include <ranges>
#include <sstream>

namespace tools {
  template <::std::ranges::range R, typename T>
  ::std::string join(R&& e, const T& d) {
    ::std::ostringstream ss;
    auto it = ::std::ranges::begin(e);
    const auto end = ::std::ranges::end(e);
    if (it != end) {
      ss << *it;
      for (++it; it != end; ++it) {
        ss << d << *it;
      }
    }
    return ss.str();
  }
}


#line 1 "tools/minimum_steiner_tree.hpp"



#include <algorithm>
#include <cassert>
#include <limits>
#include <queue>
#line 9 "tools/minimum_steiner_tree.hpp"
#include <stack>
#include <tuple>
#include <type_traits>
#include <utility>
#line 1 "tools/chmin.hpp"



#line 6 "tools/chmin.hpp"

namespace tools {

  template <typename M, typename N>
  bool chmin(M& lhs, const N& rhs) {
    bool updated;
    if constexpr (::std::is_integral_v<M> && ::std::is_integral_v<N>) {
      updated = ::std::cmp_less(rhs, lhs);
    } else {
      updated = rhs < lhs;
    }
    if (updated) lhs = rhs;
    return updated;
  }
}


#line 1 "tools/greater_by_second.hpp"



#line 5 "tools/greater_by_second.hpp"

namespace tools {

  class greater_by_second {
  public:
    template <class T1, class T2>
    bool operator()(const ::std::pair<T1, T2>& x, const ::std::pair<T1, T2>& y) const {
      return x.second > y.second;
    }
  };
}


#line 16 "tools/minimum_steiner_tree.hpp"

namespace tools {
  template <typename Cost>
  class minimum_steiner_tree {
  public:
    struct edge {
      int from;
      int to;
      Cost cost;
    };

  private:
    ::std::vector<edge> m_edges;
    ::std::vector<::std::vector<int>> m_graph;

  public:
    struct query_result {
      Cost cost;
      ::std::vector<int> vertices;
      ::std::vector<int> edge_ids;
    };

    minimum_steiner_tree() = default;
    explicit minimum_steiner_tree(const int n) : m_graph(n) {
      assert(n >= 0);
    }

    int size() const {
      return this->m_graph.size();
    }
    int add_edge(int u, int v, const Cost w) {
      assert(0 <= u && u < this->size());
      assert(0 <= v && v < this->size());
      assert(w >= 0);

      ::std::tie(u, v) = ::std::minmax({u, v});
      this->m_edges.push_back({u, v, w});
      this->m_graph[u].push_back(this->m_edges.size() - 1);
      this->m_graph[v].push_back(this->m_edges.size() - 1);
      return this->m_edges.size() - 1;
    }
    const edge& get_edge(const int i) const {
      assert(0 <= i && ::std::cmp_less(i, this->m_edges.size()));
      return this->m_edges[i];
    }
    const ::std::vector<edge>& edges() const {
      return this->m_edges;
    }

  public:
    template <bool Restore, ::std::ranges::range R>
    ::std::conditional_t<Restore, query_result, Cost> query(R&& S) const {
      if constexpr (::std::ranges::forward_range<R>) {
        const auto k = ::std::ranges::distance(S);
        assert(0 <= k && k <= this->size());
        assert(::std::ranges::all_of(S, [this](const int v) { return 0 <= v && v < this->size(); }));

        if (k == 0) {
          if constexpr (Restore) {
            return query_result{};
          } else {
            return 0;
          }
        }

        ::std::vector dp(1 << k, ::std::vector(this->size(), ::std::numeric_limits<Cost>::max()));
        ::std::vector prev(Restore ? 1 << k : 0, ::std::vector(this->size(), -1));
        for (int t = 0; const auto v : S) {
          dp[1 << t][v] = 0;
          ++t;
        }
        for (int T = 1; T < 1 << k; ++T) {
          for (int v = 0; v < this->size(); ++v) {
            for (int U = (T - 1) & T; U > 0; U = (U - 1) & T) {
              if (dp[U][v] < ::std::numeric_limits<Cost>::max() && dp[T ^ U][v] < ::std::numeric_limits<Cost>::max() && ::tools::chmin(dp[T][v], dp[U][v] + dp[T ^ U][v])) {
                if constexpr (Restore) {
                  prev[T][v] = U;
                }
              }
            }
          }
          ::std::priority_queue<::std::pair<int, Cost>, ::std::vector<::std::pair<int, Cost>>, ::tools::greater_by_second> pq;
          for (int v = 0; v < this->size(); ++v) {
            if (dp[T][v] < ::std::numeric_limits<Cost>::max()) {
              pq.emplace(v, dp[T][v]);
            }
          }
          while (!pq.empty()) {
            const auto [here, d] = pq.top();
            pq.pop();
            if (dp[T][here] < d) continue;
            for (const auto edge_id : this->m_graph[here]) {
              const auto& edge = this->m_edges[edge_id];
              const auto next = edge.from ^ edge.to ^ here;
              if (::tools::chmin(dp[T][next], dp[T][here] + edge.cost)) {
                pq.emplace(next, dp[T][next]);
                if constexpr (Restore) {
                  prev[T][next] = (1 << k) + edge_id;
                }
              }
            }
          }
        }

        if constexpr (Restore) {
          query_result qr;
          qr.cost = dp.back()[*::std::ranges::begin(S)];
          if (qr.cost == ::std::numeric_limits<Cost>::max()) return qr;

          ::std::stack<::std::pair<int, int>> stack;
          stack.emplace((1 << k) - 1, *::std::ranges::begin(S));
          qr.vertices.push_back(*::std::ranges::begin(S));
          while (!stack.empty()) {
            const auto [T, v] = stack.top();
            stack.pop();
            if (prev[T][v] == -1) continue;
            if (prev[T][v] < 1 << k) {
              stack.emplace(prev[T][v], v);
              stack.emplace(T ^ prev[T][v], v);
            } else {
              const auto edge_id = prev[T][v] - (1 << k);
              const auto& edge = this->m_edges[edge_id];
              qr.vertices.push_back(edge.from ^ edge.to ^ v);
              qr.edge_ids.push_back(edge_id);
              stack.emplace(T, edge.from ^ edge.to ^ v);
            }
          }
          return qr;
        } else {
          return dp.back()[*::std::ranges::begin(S)];
        }
      } else {
        return this->query<Restore>(::std::vector<::std::decay_t<::std::ranges::range_value_t<R>>>(::std::ranges::begin(S), ::std::ranges::end(S)));
      }
    }
    template <::std::ranges::range R>
    auto query(R&& S) const {
      return this->query<false>(::std::forward<R>(S));
    }
  };
}


#line 7 "tests/minimum_steiner_tree/restore.test.cpp"

using ll = long long;

int main() {
  std::cin.tie(nullptr);
  std::ios_base::sync_with_stdio(false);

  int N, M;
  std::cin >> N >> M;
  tools::minimum_steiner_tree<ll> graph(N);
  for (int i = 0; i < M; ++i) {
    int U, V, W;
    std::cin >> U >> V >> W;
    graph.add_edge(U, V, W);
  }
  int K;
  std::cin >> K;
  std::vector<int> X(K);
  for (auto& X_i : X) std::cin >> X_i;

  const auto qr = graph.query<true>(X);
  std::cout << qr.cost << ' ' << qr.edge_ids.size() << '\n';
  std::cout << tools::join(qr.edge_ids, " ") << '\n';

  return 0;
}

Test cases

Env	Name	Status	Elapsed	Memory
g++	example_00	AC	5 ms	4 MB
g++	example_01	AC	4 ms	4 MB
g++	example_02	AC	4 ms	4 MB
g++	overflow_killer_00	AC	164 ms	8 MB
g++	random_00	AC	5 ms	4 MB
g++	random_01	AC	5 ms	4 MB
g++	random_02	AC	6 ms	4 MB
g++	random_03	AC	4 ms	4 MB
g++	random_04	AC	4 ms	4 MB
g++	random_complete_00	AC	80 ms	6 MB
g++	random_complete_01	AC	81 ms	6 MB
g++	random_complete_02	AC	79 ms	6 MB
g++	random_complete_03	AC	80 ms	6 MB
g++	random_complete_04	AC	80 ms	6 MB
g++	random_max_00	AC	172 ms	8 MB
g++	random_max_01	AC	172 ms	8 MB
g++	random_max_02	AC	171 ms	8 MB
g++	random_max_03	AC	178 ms	8 MB
g++	random_max_04	AC	177 ms	8 MB
g++	random_small_00	AC	4 ms	4 MB
g++	random_small_01	AC	4 ms	4 MB
g++	random_small_02	AC	10 ms	4 MB
g++	random_small_03	AC	4 ms	3 MB
g++	random_small_04	AC	4 ms	4 MB