Skip to content

k_shortest_walk.hpp

SECTIONGraph INCLUDEnoya/k_shortest_walk.hpp

按长度顺序求源点到终点的前 \(k\) 条游走,允许重复经过顶点和边。

Complexity: Time: O((V + E) log V + k log k). Space: O(V + E log V + k).

AC 记录:k_shortest_walk

跳到代码 · GitHub ↗

Implementation

当前头文件,省略 include guard;依赖见 #include

/// @complexity Time: O((V + E) log V + k log k).
/// Space: O(V + E log V + k).

#include <algorithm>
#include <cassert>
#include <functional>
#include <limits>
#include <queue>
#include <tuple>
#include <utility>
#include <vector>

namespace noya {
namespace k_shortest_walk_internal {

template <class Weight> struct persistent_leftist_heap {
  struct node {
    Weight val{};
    int to1 = -1;
    int l = -1;
    int r = -1;
    int rk = 1;
  };

  std::vector<node> nd;

  int node_rank(int id) const { return id == -1 ? 0 : nd[id].rk; }

  int meld(int a, int b) {
    if (a == -1 || b == -1) {
      return a == -1 ? b : a;
    }
    if (nd[b].val < nd[a].val) {
      std::swap(a, b);
    }
    node cp = nd[a];
    cp.r = meld(cp.r, b);
    if (node_rank(cp.l) < node_rank(cp.r)) {
      std::swap(cp.l, cp.r);
    }
    cp.rk = node_rank(cp.r) + 1;
    nd.push_back(cp);
    return int(nd.size()) - 1;
  }

  int insert(int rt, Weight val, int to1) {
    nd.push_back(node{val, to1});
    return meld(rt, int(nd.size()) - 1);
  }
};

} // namespace k_shortest_walk_internal

/// @brief Return the lengths of the first cnt directed walks from s to t.
/// Reverse Dijkstra fixes one shortest-path tree toward the target. Every
/// non-tree edge has a nonnegative sidetrack cost
/// `w + dis[v] - dis[u]` for an edge (u, v, w); a walk is its shortest base path plus an
/// ordered sequence of sidetracks. Persistent leftist heaps collect all
/// sidetracks available along each tree path. A final priority queue explores
/// replacing a chosen sidetrack by a heap child or appending one after its
/// endpoint, so walks are emitted in nondecreasing total length. Vertices
/// and edges may repeat, including through zero-weight cycles.
template <class Weight>
std::vector<Weight>
k_shortest_walk_lengths(int n,
                        const std::vector<std::tuple<int, int, Weight>> &es,
                        int s, int t, int cnt) {
  assert(n >= 0);
  assert(0 <= s && s < n);
  assert(0 <= t && t < n);
  assert(cnt >= 0);
  if (cnt == 0) {
    return {};
  }

  struct adjacent_edge {
    int u;
    Weight w;
    int id;
  };
  std::vector<std::vector<adjacent_edge>> out(n);
  std::vector<std::vector<adjacent_edge>> rev(n);
  for (int id = 0; id < int(es.size()); id++) {
    auto [u1, to, w] = es[id];
    assert(0 <= u1 && u1 < n);
    assert(0 <= to && to < n);
    assert(w >= Weight{});
    out[u1].push_back({to, w, id});
    rev[to].push_back({u1, w, id});
  }

  const Weight inf = std::numeric_limits<Weight>::max() / 4;
  std::vector<Weight> dis(n, inf);
  std::vector<int> pe(n, -1);
  std::vector<int> par(n, -1);
  using distance_entry = std::pair<Weight, int>;
  std::priority_queue<distance_entry, std::vector<distance_entry>,
                      std::greater<distance_entry>>
      pq;
  dis[t] = Weight{};
  pq.emplace(Weight{}, t);
  while (!pq.empty()) {
    auto [du, u] = pq.top();
    pq.pop();
    if (du != dis[u]) {
      continue;
    }
    for (auto [pre, w, id] : rev[u]) {
      Weight can = du + w;
      if (can < dis[pre]) {
        dis[pre] = can;
        par[pre] = u;
        pe[pre] = id;
        pq.emplace(can, pre);
      }
    }
  }
  if (dis[s] == inf) {
    return {};
  }

  std::vector<std::vector<int>> ch(n);
  for (int u = 0; u < n; u++) {
    if (par[u] != -1) {
      ch[par[u]].push_back(u);
    }
  }

  using heap_type = k_shortest_walk_internal::persistent_leftist_heap<Weight>;
  heap_type hp;
  std::vector<int> rot(n, -1);
  std::vector<int> ord{t};
  for (int i = 0; i < int(ord.size()); i++) {
    int u = ord[i];
    if (par[u] != -1) {
      rot[u] = rot[par[u]];
    }
    for (auto [to, w, id] : out[u]) {
      if (id == pe[u] || dis[to] == inf) {
        continue;
      }
      Weight sid = w + dis[to] - dis[u];
      rot[u] = hp.insert(rot[u], sid, to);
    }
    ord.insert(ord.end(), ch[u].begin(), ch[u].end());
  }

  using candidate = std::pair<Weight, int>;
  std::priority_queue<candidate, std::vector<candidate>,
                      std::greater<candidate>>
      ca1;
  ca1.emplace(dis[s], -1);
  std::vector<Weight> res;
  res.reserve(cnt);
  while (!ca1.empty() && int(res.size()) < cnt) {
    auto [tot, ni] = ca1.top();
    ca1.pop();
    res.push_back(tot);

    int cr;
    if (ni == -1) {
      cr = rot[s];
    } else {
      const auto &cur = hp.nd[ni];
      for (int v : {cur.l, cur.r}) {
        if (v != -1) {
          ca1.emplace(tot + hp.nd[v].val - cur.val, v);
        }
      }
      cr = rot[cur.to1];
    }
    if (cr != -1) {
      ca1.emplace(tot + hp.nd[cr].val, cr);
    }
  }
  return res;
}

} // namespace noya
#ifndef NOYA_K_SHORTEST_WALK_HPP
#define NOYA_K_SHORTEST_WALK_HPP 1

/// @complexity Time: O((V + E) log V + k log k).
/// Space: O(V + E log V + k).

#include <algorithm>
#include <cassert>
#include <functional>
#include <limits>
#include <queue>
#include <tuple>
#include <utility>
#include <vector>

namespace noya {
namespace k_shortest_walk_internal {

template <class Weight> struct persistent_leftist_heap {
  struct node {
    Weight val{};
    int to1 = -1;
    int l = -1;
    int r = -1;
    int rk = 1;
  };

  std::vector<node> nd;

  int node_rank(int id) const { return id == -1 ? 0 : nd[id].rk; }

  int meld(int a, int b) {
    if (a == -1 || b == -1) {
      return a == -1 ? b : a;
    }
    if (nd[b].val < nd[a].val) {
      std::swap(a, b);
    }
    node cp = nd[a];
    cp.r = meld(cp.r, b);
    if (node_rank(cp.l) < node_rank(cp.r)) {
      std::swap(cp.l, cp.r);
    }
    cp.rk = node_rank(cp.r) + 1;
    nd.push_back(cp);
    return int(nd.size()) - 1;
  }

  int insert(int rt, Weight val, int to1) {
    nd.push_back(node{val, to1});
    return meld(rt, int(nd.size()) - 1);
  }
};

} // namespace k_shortest_walk_internal

/// @brief Return the lengths of the first cnt directed walks from s to t.
/// Reverse Dijkstra fixes one shortest-path tree toward the target. Every
/// non-tree edge has a nonnegative sidetrack cost
/// `w + dis[v] - dis[u]` for an edge (u, v, w); a walk is its shortest base path plus an
/// ordered sequence of sidetracks. Persistent leftist heaps collect all
/// sidetracks available along each tree path. A final priority queue explores
/// replacing a chosen sidetrack by a heap child or appending one after its
/// endpoint, so walks are emitted in nondecreasing total length. Vertices
/// and edges may repeat, including through zero-weight cycles.
template <class Weight>
std::vector<Weight>
k_shortest_walk_lengths(int n,
                        const std::vector<std::tuple<int, int, Weight>> &es,
                        int s, int t, int cnt) {
  assert(n >= 0);
  assert(0 <= s && s < n);
  assert(0 <= t && t < n);
  assert(cnt >= 0);
  if (cnt == 0) {
    return {};
  }

  struct adjacent_edge {
    int u;
    Weight w;
    int id;
  };
  std::vector<std::vector<adjacent_edge>> out(n);
  std::vector<std::vector<adjacent_edge>> rev(n);
  for (int id = 0; id < int(es.size()); id++) {
    auto [u1, to, w] = es[id];
    assert(0 <= u1 && u1 < n);
    assert(0 <= to && to < n);
    assert(w >= Weight{});
    out[u1].push_back({to, w, id});
    rev[to].push_back({u1, w, id});
  }

  const Weight inf = std::numeric_limits<Weight>::max() / 4;
  std::vector<Weight> dis(n, inf);
  std::vector<int> pe(n, -1);
  std::vector<int> par(n, -1);
  using distance_entry = std::pair<Weight, int>;
  std::priority_queue<distance_entry, std::vector<distance_entry>,
                      std::greater<distance_entry>>
      pq;
  dis[t] = Weight{};
  pq.emplace(Weight{}, t);
  while (!pq.empty()) {
    auto [du, u] = pq.top();
    pq.pop();
    if (du != dis[u]) {
      continue;
    }
    for (auto [pre, w, id] : rev[u]) {
      Weight can = du + w;
      if (can < dis[pre]) {
        dis[pre] = can;
        par[pre] = u;
        pe[pre] = id;
        pq.emplace(can, pre);
      }
    }
  }
  if (dis[s] == inf) {
    return {};
  }

  std::vector<std::vector<int>> ch(n);
  for (int u = 0; u < n; u++) {
    if (par[u] != -1) {
      ch[par[u]].push_back(u);
    }
  }

  using heap_type = k_shortest_walk_internal::persistent_leftist_heap<Weight>;
  heap_type hp;
  std::vector<int> rot(n, -1);
  std::vector<int> ord{t};
  for (int i = 0; i < int(ord.size()); i++) {
    int u = ord[i];
    if (par[u] != -1) {
      rot[u] = rot[par[u]];
    }
    for (auto [to, w, id] : out[u]) {
      if (id == pe[u] || dis[to] == inf) {
        continue;
      }
      Weight sid = w + dis[to] - dis[u];
      rot[u] = hp.insert(rot[u], sid, to);
    }
    ord.insert(ord.end(), ch[u].begin(), ch[u].end());
  }

  using candidate = std::pair<Weight, int>;
  std::priority_queue<candidate, std::vector<candidate>,
                      std::greater<candidate>>
      ca1;
  ca1.emplace(dis[s], -1);
  std::vector<Weight> res;
  res.reserve(cnt);
  while (!ca1.empty() && int(res.size()) < cnt) {
    auto [tot, ni] = ca1.top();
    ca1.pop();
    res.push_back(tot);

    int cr;
    if (ni == -1) {
      cr = rot[s];
    } else {
      const auto &cur = hp.nd[ni];
      for (int v : {cur.l, cur.r}) {
        if (v != -1) {
          ca1.emplace(tot + hp.nd[v].val - cur.val, v);
        }
      }
      cr = rot[cur.to1];
    }
    if (cr != -1) {
      ca1.emplace(tot + hp.nd[cr].val, cr);
    }
  }
  return res;
}

} // namespace noya

#endif // NOYA_K_SHORTEST_WALK_HPP
#include <algorithm>
#include <cassert>
#include <functional>
#include <limits>
#include <queue>
#include <tuple>
#include <utility>
#include <vector>

/// @complexity Time: O((V + E) log V + k log k).
/// Space: O(V + E log V + k).

namespace noya {
namespace k_shortest_walk_internal {

template <class Weight> struct persistent_leftist_heap {
  struct node {
    Weight val{};
    int to1 = -1;
    int l = -1;
    int r = -1;
    int rk = 1;
  };

  std::vector<node> nd;

  int node_rank(int id) const { return id == -1 ? 0 : nd[id].rk; }

  int meld(int a, int b) {
    if (a == -1 || b == -1) {
      return a == -1 ? b : a;
    }
    if (nd[b].val < nd[a].val) {
      std::swap(a, b);
    }
    node cp = nd[a];
    cp.r = meld(cp.r, b);
    if (node_rank(cp.l) < node_rank(cp.r)) {
      std::swap(cp.l, cp.r);
    }
    cp.rk = node_rank(cp.r) + 1;
    nd.push_back(cp);
    return int(nd.size()) - 1;
  }

  int insert(int rt, Weight val, int to1) {
    nd.push_back(node{val, to1});
    return meld(rt, int(nd.size()) - 1);
  }
};

} // namespace k_shortest_walk_internal

/// @brief Return the lengths of the first cnt directed walks from s to t.
/// Reverse Dijkstra fixes one shortest-path tree toward the target. Every
/// non-tree edge has a nonnegative sidetrack cost
/// `w + dis[v] - dis[u]` for an edge (u, v, w); a walk is its shortest base path plus an
/// ordered sequence of sidetracks. Persistent leftist heaps collect all
/// sidetracks available along each tree path. A final priority queue explores
/// replacing a chosen sidetrack by a heap child or appending one after its
/// endpoint, so walks are emitted in nondecreasing total length. Vertices
/// and edges may repeat, including through zero-weight cycles.
template <class Weight>
std::vector<Weight>
k_shortest_walk_lengths(int n,
                        const std::vector<std::tuple<int, int, Weight>> &es,
                        int s, int t, int cnt) {
  assert(n >= 0);
  assert(0 <= s && s < n);
  assert(0 <= t && t < n);
  assert(cnt >= 0);
  if (cnt == 0) {
    return {};
  }

  struct adjacent_edge {
    int u;
    Weight w;
    int id;
  };
  std::vector<std::vector<adjacent_edge>> out(n);
  std::vector<std::vector<adjacent_edge>> rev(n);
  for (int id = 0; id < int(es.size()); id++) {
    auto [u1, to, w] = es[id];
    assert(0 <= u1 && u1 < n);
    assert(0 <= to && to < n);
    assert(w >= Weight{});
    out[u1].push_back({to, w, id});
    rev[to].push_back({u1, w, id});
  }

  const Weight inf = std::numeric_limits<Weight>::max() / 4;
  std::vector<Weight> dis(n, inf);
  std::vector<int> pe(n, -1);
  std::vector<int> par(n, -1);
  using distance_entry = std::pair<Weight, int>;
  std::priority_queue<distance_entry, std::vector<distance_entry>,
                      std::greater<distance_entry>>
      pq;
  dis[t] = Weight{};
  pq.emplace(Weight{}, t);
  while (!pq.empty()) {
    auto [du, u] = pq.top();
    pq.pop();
    if (du != dis[u]) {
      continue;
    }
    for (auto [pre, w, id] : rev[u]) {
      Weight can = du + w;
      if (can < dis[pre]) {
        dis[pre] = can;
        par[pre] = u;
        pe[pre] = id;
        pq.emplace(can, pre);
      }
    }
  }
  if (dis[s] == inf) {
    return {};
  }

  std::vector<std::vector<int>> ch(n);
  for (int u = 0; u < n; u++) {
    if (par[u] != -1) {
      ch[par[u]].push_back(u);
    }
  }

  using heap_type = k_shortest_walk_internal::persistent_leftist_heap<Weight>;
  heap_type hp;
  std::vector<int> rot(n, -1);
  std::vector<int> ord{t};
  for (int i = 0; i < int(ord.size()); i++) {
    int u = ord[i];
    if (par[u] != -1) {
      rot[u] = rot[par[u]];
    }
    for (auto [to, w, id] : out[u]) {
      if (id == pe[u] || dis[to] == inf) {
        continue;
      }
      Weight sid = w + dis[to] - dis[u];
      rot[u] = hp.insert(rot[u], sid, to);
    }
    ord.insert(ord.end(), ch[u].begin(), ch[u].end());
  }

  using candidate = std::pair<Weight, int>;
  std::priority_queue<candidate, std::vector<candidate>,
                      std::greater<candidate>>
      ca1;
  ca1.emplace(dis[s], -1);
  std::vector<Weight> res;
  res.reserve(cnt);
  while (!ca1.empty() && int(res.size()) < cnt) {
    auto [tot, ni] = ca1.top();
    ca1.pop();
    res.push_back(tot);

    int cr;
    if (ni == -1) {
      cr = rot[s];
    } else {
      const auto &cur = hp.nd[ni];
      for (int v : {cur.l, cur.r}) {
        if (v != -1) {
          ca1.emplace(tot + hp.nd[v].val - cur.val, v);
        }
      }
      cr = rot[cur.to1];
    }
    if (cr != -1) {
      ca1.emplace(tot + hp.nd[cr].val, cr);
    }
  }
  return res;
}

} // namespace noya