Add "array" parameter to {dfs,bfs,dijkstra,astar}_iterator()

src/graph/numpy_bind.hh

@@ -65,7 +65,7 @@ typedef boost::mpl::map<
 template <class ValueType>
 boost::python::object wrap_vector_owned(const std::vector<ValueType>& vec)
 {
-    int val_type = boost::mpl::at<numpy_types,ValueType>::type::value;
+    size_t val_type = boost::mpl::at<numpy_types,ValueType>::type::value;
     npy_intp size[1];
     size[0] = vec.size();
     PyArrayObject* ndarray;
@@ -91,13 +91,12 @@ template <class ValueType>
 boost::python::object wrap_vector_not_owned(std::vector<ValueType>& vec)
 {
     PyArrayObject* ndarray;
-    int val_type = boost::mpl::at<numpy_types,ValueType>::type::value;
+    size_t val_type = boost::mpl::at<numpy_types,ValueType>::type::value;
     npy_intp size = vec.size();
     if (vec.empty())
         return wrap_vector_owned(vec); // return an _owned_ array of size one.
-    else
-        ndarray = (PyArrayObject*) PyArray_SimpleNewFromData(1, &size, val_type,
-                                                             vec.data());
+    ndarray = (PyArrayObject*) PyArray_SimpleNewFromData(1, &size, val_type,
+                                                         vec.data());
     PyArray_ENABLEFLAGS(ndarray,NPY_ARRAY_ALIGNED | NPY_ARRAY_C_CONTIGUOUS |
                         NPY_ARRAY_WRITEABLE);
     boost::python::handle<> x((PyObject*) ndarray);
@@ -105,16 +104,60 @@ boost::python::object wrap_vector_not_owned(std::vector<ValueType>& vec)
     return o;
 }
 
+template <class ValueType, size_t Dim>
+boost::python::object wrap_vector_owned(const std::vector<std::array<ValueType, Dim>>& vec)
+{
+    size_t n = vec.size() * Dim;
+    size_t val_type = boost::mpl::at<numpy_types,ValueType>::type::value;
+    PyArrayObject* ndarray;
+    if (n == 0)
+    {
+        npy_intp size[1];
+        size[0] = vec.size();
+        ndarray = (PyArrayObject*) PyArray_SimpleNew(1, size, val_type);
+    }
+    else
+    {
+        ValueType* new_data = new ValueType[n];
+        memcpy(new_data, vec.data(), n * sizeof(ValueType));
+        npy_intp shape[2] = {int(vec.size()), int(Dim)};
+        ndarray = (PyArrayObject*) PyArray_SimpleNewFromData(Dim, shape,
+                                                             val_type,
+                                                             new_data);
+    }
+    PyArray_ENABLEFLAGS(ndarray, NPY_ARRAY_ALIGNED | NPY_ARRAY_F_CONTIGUOUS |
+                        NPY_ARRAY_OWNDATA | NPY_ARRAY_WRITEABLE);
+    boost::python::handle<> x((PyObject*) ndarray);
+    boost::python::object o(x);
+    return o;
+}
+
+template <class ValueType, size_t Dim>
+boost::python::object wrap_vector_not_owned(const std::vector<std::array<ValueType, Dim>>& vec)
+{
+    size_t val_type = boost::mpl::at<numpy_types,ValueType>::type::value;
+    if (vec.empty())
+        return wrap_vector_owned(vec); // return an _owned_ array of size one.
+    npy_intp shape[2] = {int(vec.size()), int(Dim)};
+    PyArrayObject* ndarray =
+        (PyArrayObject*) PyArray_SimpleNewFromData(Dim, shape, val_type,
+                                                   vec.data());
+    PyArray_ENABLEFLAGS(ndarray, NPY_ARRAY_ALIGNED | NPY_ARRAY_F_CONTIGUOUS |
+                        NPY_ARRAY_WRITEABLE);
+    boost::python::handle<> x((PyObject*) ndarray);
+    boost::python::object o(x);
+    return o;
+}
 
-template <class ValueType, int Dim>
+template <class ValueType, size_t Dim>
 boost::python::object
 wrap_multi_array_owned(const boost::multi_array<ValueType,Dim>& array)
 {
     ValueType* new_data = new ValueType[array.num_elements()];
     memcpy(new_data, array.data(), array.num_elements() * sizeof(ValueType));
-    int val_type = boost::mpl::at<numpy_types,ValueType>::type::value;
+    size_t val_type = boost::mpl::at<numpy_types,ValueType>::type::value;
     npy_intp shape[Dim];
-    for (int i = 0; i < Dim; ++i)
+    for (size_t i = 0; i < Dim; ++i)
         shape[i] = array.shape()[i];
     PyArrayObject* ndarray =
         (PyArrayObject*) PyArray_SimpleNewFromData(Dim, shape, val_type,
@@ -126,11 +169,11 @@ wrap_multi_array_owned(const boost::multi_array<ValueType,Dim>& array)
     return o;
 }
 
-template <class ValueType, int Dim>
+template <class ValueType, size_t Dim>
 boost::python::object
 wrap_multi_array_not_owned(boost::multi_array<ValueType,Dim>& array)
 {
-    int val_type = boost::mpl::at<numpy_types,ValueType>::type::value;
+    size_t val_type = boost::mpl::at<numpy_types,ValueType>::type::value;
     PyArrayObject* ndarray =
         (PyArrayObject*) PyArray_SimpleNewFromData(Dim, array.shape(), val_type,
                                                    array.origin());

src/graph/search/graph_astar.cc

@@ -177,6 +177,61 @@ boost::python::object astar_search_generator_fast(GraphInterface& g,
 #endif
 }
 
+class AStarArrayVisitor: public astar_visitor<>
+{
+public:
+    AStarArrayVisitor(std::vector<std::array<size_t, 2>>& edges)
+        : _edges(edges) {}
+
+    template <class Edge, class Graph>
+    void edge_relaxed(const Edge& e, Graph& g)
+    {
+        _edges.push_back({source(e, g), target(e,g)});
+    }
+
+private:
+    std::vector<std::array<size_t, 2>>& _edges;
+};
+
+boost::python::object astar_search_array(GraphInterface& g,
+                                         size_t source,
+                                         boost::any dist_map,
+                                         boost::any weight,
+                                         python::object cmp,
+                                         python::object cmb,
+                                         python::object zero,
+                                         python::object inf,
+                                         python::object h)
+{
+    std::vector<std::array<size_t, 2>> edges;
+    AStarArrayVisitor vis(edges);
+    run_action<graph_tool::all_graph_views,mpl::true_>()
+        (g, std::bind(do_astar_search(),  std::placeholders::_1, source,
+                      std::placeholders::_2, dummy_property_map(), weight,
+                      vis, make_pair(AStarCmp(cmp), AStarCmb(cmb)),
+                      make_pair(zero, inf), h, std::ref(g)),
+         writable_vertex_properties())(dist_map);
+    return wrap_vector_owned<size_t,2>(edges);
+}
+
+boost::python::object astar_search_array_fast(GraphInterface& g,
+                                              size_t source,
+                                              boost::any dist_map,
+                                              boost::any weight,
+                                              python::object zero,
+                                              python::object inf,
+                                              python::object h)
+{
+    std::vector<std::array<size_t, 2>> edges;
+    AStarArrayVisitor vis(edges);
+    run_action<graph_tool::all_graph_views,mpl::true_>()
+        (g, std::bind(do_astar_search_fast(),  std::placeholders::_1, source,
+                      std::placeholders::_2, std::placeholders::_3,
+                      vis, make_pair(zero, inf), h, std::ref(g)),
+         writable_vertex_scalar_properties(),
+         edge_scalar_properties())(dist_map, weight);
+    return wrap_vector_owned<size_t,2>(edges);
+}
 
 void export_astar()
 {
@@ -184,4 +239,6 @@ void export_astar()
     def("astar_search", &a_star_search);
     def("astar_generator", &astar_search_generator);
     def("astar_generator_fast", &astar_search_generator_fast);
+    def("astar_array", &astar_search_array);
+    def("astar_array_fast", &astar_search_array_fast);
 }

src/graph/search/graph_bfs.cc

@@ -106,28 +106,32 @@ private:
     boost::python::object _vis;
 };
 
+template <class Graph, class Visitor>
+void do_bfs(Graph& g, size_t s, Visitor&& vis)
+{
+    typename vprop_map_t<default_color_type>::type
+        color(get(vertex_index_t(), g));
+
+    auto v = vertex(s, g);
+    if (v == graph_traits<Graph>::null_vertex())
+    {
+        for (auto u : vertices_range(g))
+        {
+            if (color[u] == color_traits<default_color_type>::black())
+                continue;
+            breadth_first_visit(g, u, visitor(vis).color_map(color));
+        }
+    }
+    else
+    {
+        breadth_first_visit(g, v, visitor(vis).color_map(color));
+    }
+}
 
 void bfs_search(GraphInterface& gi, size_t s, python::object vis)
 {
     run_action<graph_tool::all_graph_views,mpl::true_>()
-        (gi,
-         [&](auto &g)
-         {
-             typedef typename std::remove_reference<decltype(g)>::type g_t;
-             typename vprop_map_t<default_color_type>::type
-                 color(get(vertex_index_t(), g));
-             auto visw = BFSVisitorWrapper(gi, vis);
-             auto v = vertex(s, g);
-             if (v == graph_traits<g_t>::null_vertex())
-             {
-                 for (auto u : vertices_range(g))
-                     breadth_first_search(g, u, visitor(visw).color_map(color));
-             }
-             else
-             {
-                 breadth_first_visit(g, v, visitor(visw).color_map(color));
-             }
-         })();
+        (gi, [&](auto &g){ do_bfs(g, s, BFSVisitorWrapper(gi, vis)); })();
 }
 
 #ifdef HAVE_BOOST_COROUTINE
@@ -160,28 +164,7 @@ boost::python::object bfs_search_generator(GraphInterface& g, size_t s)
         {
             BFSGeneratorVisitor vis(g, yield);
             run_action<graph_tool::all_graph_views,mpl::true_>()
-                (g,
-                 [&](auto &g)
-                 {
-                     typedef typename std::remove_reference<decltype(g)>::type g_t;
-                     typename vprop_map_t<default_color_type>::type
-                         color(get(vertex_index_t(), g));
-
-                     auto v = vertex(s, g);
-                     if (v == graph_traits<g_t>::null_vertex())
-                     {
-                         for (auto u : vertices_range(g))
-                         {
-                             if (color[u] == color_traits<default_color_type>::black())
-                                 continue;
-                             breadth_first_visit(g, u, visitor(vis).color_map(color));
-                         }
-                     }
-                     else
-                     {
-                         breadth_first_visit(g, v, visitor(vis).color_map(color));
-                     }
-                 })();
+                (g, [&](auto &g){ do_bfs(g, s, vis); })();
         };
     return boost::python::object(CoroGenerator(dispatch));
 #else
@@ -189,9 +172,36 @@ boost::python::object bfs_search_generator(GraphInterface& g, size_t s)
 #endif
 }
 
+class BFSArrayVisitor : public bfs_visitor<>
+{
+public:
+    BFSArrayVisitor(std::vector<std::array<size_t, 2>>& edges)
+        : _edges(edges) {}
+
+    template <class Edge, class Graph>
+    void tree_edge(const Edge& e, Graph& g)
+    {
+        _edges.push_back({source(e, g), target(e,g)});
+    }
+
+private:
+    std::vector<std::array<size_t, 2>>& _edges;
+};
+
+boost::python::object bfs_search_array(GraphInterface& g, size_t s)
+{
+    std::vector<std::array<size_t, 2>> edges;
+    BFSArrayVisitor vis(edges);
+    run_action<graph_tool::all_graph_views,mpl::true_>()
+        (g, [&](auto &g){ do_bfs(g, s, vis); })();
+    return wrap_vector_owned<size_t,2>(edges);
+}
+
+
 void export_bfs()
 {
     using namespace boost::python;
     def("bfs_search", &bfs_search);
     def("bfs_search_generator", &bfs_search_generator);
+    def("bfs_search_array", &bfs_search_array);
 }

src/graph/search/graph_dfs.cc

@@ -90,22 +90,22 @@ private:
     python::object _vis;
 };
 
+template <class Graph, class Visitor>
+void do_dfs(Graph& g, size_t s, Visitor&& vis)
+{
+    typename vprop_map_t<default_color_type>::type
+        color(get(vertex_index_t(), g));
+    auto v = vertex(s, g);
+    if (v == graph_traits<Graph>::null_vertex())
+        depth_first_search(g, vis, color);
+    else
+        depth_first_visit(g, v, vis, color);
+}
+
 void dfs_search(GraphInterface& gi, size_t s, python::object vis)
 {
     run_action<graph_tool::all_graph_views, mpl::true_>()
-        (gi,
-         [&](auto &g)
-         {
-             typedef typename std::remove_reference<decltype(g)>::type g_t;
-             typename vprop_map_t<default_color_type>::type
-                 color(get(vertex_index_t(), g));
-             auto visw = DFSVisitorWrapper(gi, vis);
-             auto v = vertex(s, g);
-             if (v == graph_traits<g_t>::null_vertex())
-                 depth_first_search(g, visw, color);
-             else
-                 depth_first_visit(g, v, visw, color);
-         })();
+        (gi, [&](auto &g) { do_dfs(g, s, DFSVisitorWrapper(gi, vis));})();
 }
 
 #ifdef HAVE_BOOST_COROUTINE
@@ -139,18 +139,7 @@ boost::python::object dfs_search_generator(GraphInterface& g, size_t s)
         {
             DFSGeneratorVisitor vis(g, yield);
             run_action<graph_tool::all_graph_views,mpl::true_>()
-                (g,
-                 [&](auto &g)
-                 {
-                     typedef typename std::remove_reference<decltype(g)>::type g_t;
-                     typename vprop_map_t<default_color_type>::type
-                         color(get(vertex_index_t(), g));
-                     auto v = vertex(s, g);
-                     if (v == graph_traits<g_t>::null_vertex())
-                         depth_first_search(g, vis, color);
-                     else
-                         depth_first_visit(g, v, vis, color);
-                 })();
+                (g, [&](auto &g) { do_dfs(g, s, vis);})();
         };
     return boost::python::object(CoroGenerator(dispatch));
 #else
@@ -158,9 +147,35 @@ boost::python::object dfs_search_generator(GraphInterface& g, size_t s)
 #endif
 }
 
+class DFSArrayVisitor: public dfs_visitor<>
+{
+public:
+    DFSArrayVisitor(std::vector<std::array<size_t, 2>>& edges)
+        : _edges(edges) {}
+
+    template <class Edge, class Graph>
+    void tree_edge(const Edge& e, Graph& g)
+    {
+        _edges.push_back({source(e, g), target(e,g)});
+    }
+
+private:
+    std::vector<std::array<size_t, 2>>& _edges;
+};
+
+boost::python::object dfs_search_array(GraphInterface& g, size_t s)
+{
+    std::vector<std::array<size_t, 2>> edges;
+    DFSArrayVisitor vis(edges);
+    run_action<graph_tool::all_graph_views,mpl::true_>()
+        (g, [&](auto &g){ do_dfs(g, s, vis); })();
+    return wrap_vector_owned<size_t,2>(edges);
+}
+
 void export_dfs()
 {
     using namespace boost::python;
     def("dfs_search", &dfs_search);
     def("dfs_search_generator", &dfs_search_generator);
+    def("dfs_search_array", &dfs_search_array);
 }

src/graph/search/graph_dijkstra.cc

@@ -297,10 +297,67 @@ boost::python::object dijkstra_search_generator_fast(GraphInterface& g,
 #endif
 }
 
+class DJKArrayVisitor: public dijkstra_visitor<>
+{
+public:
+    DJKArrayVisitor(std::vector<std::array<size_t, 2>>& edges)
+        : _edges(edges) {}
+
+    template <class Edge, class Graph>
+    void edge_relaxed(const Edge& e, Graph& g)
+    {
+        _edges.push_back({source(e, g), target(e,g)});
+    }
+
+private:
+    std::vector<std::array<size_t, 2>>& _edges;
+};
+
+
+boost::python::object dijkstra_search_array(GraphInterface& g,
+                                            size_t source,
+                                            boost::any dist_map,
+                                            boost::any weight,
+                                            python::object cmp,
+                                            python::object cmb,
+                                            python::object zero,
+                                            python::object inf)
+{
+    std::vector<std::array<size_t, 2>> edges;
+    DJKArrayVisitor vis(edges);
+    run_action<graph_tool::all_graph_views, mpl::true_>()
+        (g, std::bind(do_djk_search(), std::placeholders::_1, source,
+                      std::placeholders::_2, dummy_property_map(), weight,
+                      vis, DJKCmp(cmp), DJKCmb(cmb),
+                      make_pair(zero, inf)),
+         writable_vertex_properties())(dist_map);
+    return wrap_vector_owned<size_t,2>(edges);
+}
+
+boost::python::object dijkstra_search_array_fast(GraphInterface& g,
+                                                 size_t source,
+                                                 boost::any dist_map,
+                                                 boost::any weight,
+                                                 python::object zero,
+                                                 python::object inf)
+{
+    std::vector<std::array<size_t, 2>> edges;
+    DJKArrayVisitor vis(edges);
+    run_action<graph_tool::all_graph_views, mpl::true_>()
+        (g, std::bind(do_djk_search_fast(), std::placeholders::_1, source,
+                      std::placeholders::_2, std::placeholders::_3,
+                      vis, make_pair(zero, inf)),
+         writable_vertex_scalar_properties(),
+         edge_scalar_properties())(dist_map, weight);
+    return wrap_vector_owned<size_t,2>(edges);
+}
+
 void export_dijkstra()
 {
     using namespace boost::python;
     def("dijkstra_search", &dijkstra_search);
     def("dijkstra_generator", &dijkstra_search_generator);
     def("dijkstra_generator_fast", &dijkstra_search_generator_fast);
+    def("dijkstra_array", &dijkstra_search_array);
+    def("dijkstra_array_fast", &dijkstra_search_array_fast);
 }

src/graph_tool/search/__init__.py

@@ -295,7 +295,7 @@ def bfs_search(g, source=None, visitor=BFSVisitor()):
     except StopSearch:
         pass
 
-def bfs_iterator(g, source=None):
+def bfs_iterator(g, source=None, array=False):
     r"""Return an iterator of the edges corresponding to a breath-first traversal of
     the graph.
 
@@ -307,10 +307,16 @@ def bfs_iterator(g, source=None):
         Source vertex. If unspecified, all vertices will be traversed, by
         iterating over starting vertices according to their index in increasing
         order.
+    array : ``bool`` (optional, default: ``False``)
+        If ``True``, a :class:`numpy.ndarray` will the edge endpoints be
+        returned instead.
 
     Returns
     -------
-    bfs_iterator : An iterator over the edges in breath-first order.
+    bfs_iterator : Iterator or :class:`numpy.ndarray`
+        An iterator over the edges in breath-first order. If ``array == True``,
+        this will be a :class:`numpy.ndarray` instead, of shape ``(E,2)``,
+        containing the edge endpoints.
 
     See Also
     --------
@@ -348,12 +354,16 @@ def bfs_iterator(g, source=None):
              Symposium on the Theory of Switching, 1959
     .. [bfs-bgl] http://www.boost.org/doc/libs/release/libs/graph/doc/breadth_first_search.html
     .. [bfs-wikipedia] http://en.wikipedia.org/wiki/Breadth-first_search
+
     """
     if source is None:
         source = _get_null_vertex()
     else:
         source = int(source)
-    return libgraph_tool_search.bfs_search_generator(g._Graph__graph, source)
+    if not array:
+        return libgraph_tool_search.bfs_search_generator(g._Graph__graph, source)
+    else:
+        return libgraph_tool_search.bfs_search_array(g._Graph__graph, source)
 
 
 class DFSVisitor(object):
@@ -589,7 +599,7 @@ def dfs_search(g, source=None, visitor=DFSVisitor()):
     except StopSearch:
         pass
 
-def dfs_iterator(g, source=None):
+def dfs_iterator(g, source=None, array=False):
     r"""Return an iterator of the edges corresponding to a depth-first traversal of
     the graph.
 
@@ -601,10 +611,16 @@ def dfs_iterator(g, source=None):
         Source vertex. If unspecified, all vertices will be traversed, by
         iterating over starting vertices according to their index in increasing
         order.
+    array : ``bool`` (optional, default: ``False``)
+        If ``True``, a :class:`numpy.ndarray` will the edge endpoints be
+        returned instead.
 
     Returns
     -------
-    dfs_iterator : An iterator over the edges in detpth-first order.
+    dfs_iterator : Iterator or :class:`numpy.ndarray`
+        An iterator over the edges in depth-first order. If ``array == True``,
+        this will be a :class:`numpy.ndarray` instead, of shape ``(E,2)``,
+        containing the edge endpoints.
 
     See Also
     --------
@@ -640,13 +656,17 @@ def dfs_iterator(g, source=None):
     ----------
     .. [dfs-bgl] http://www.boost.org/doc/libs/release/libs/graph/doc/depth_first_search.html
     .. [dfs-wikipedia] http://en.wikipedia.org/wiki/Depth-first_search
+
     """
 
     if source is None:
         source = _get_null_vertex()
     else:
         source = int(source)
-    return libgraph_tool_search.dfs_search_generator(g._Graph__graph, source)
+    if not array:
+        return libgraph_tool_search.dfs_search_generator(g._Graph__graph, source)
+    else:
+        return libgraph_tool_search.dfs_search_array(g._Graph__graph, source)
 
 class DijkstraVisitor(object):
     r"""A visitor object that is invoked at the event-points inside the
@@ -952,7 +972,7 @@ def dijkstra_search(g, weight, source=None, visitor=DijkstraVisitor(), dist_map=
     return dist_map, pred_map
 
 def dijkstra_iterator(g, weight, source=None, dist_map=None, combine=None,
-                      compare=None, zero=0, infinity=numpy.inf):
+                      compare=None, zero=0, infinity=numpy.inf, array=False):
     r"""Return an iterator of the edges corresponding to a Dijkstra traversal of
     the graph.
 
@@ -979,12 +999,18 @@ def dijkstra_iterator(g, weight, source=None, dist_map=None, combine=None,
          Value assumed to correspond to a distance of zero by the combine and
          compare functions.
     infinity : int or float (optional, default: ``numpy.inf``)
-         Value assumed to correspond to a distance of infinity by the combine and
-         compare functions.
+         Value assumed to correspond to a distance of infinity by the combine
+         and compare functions.
+    array : ``bool`` (optional, default: ``False``)
+        If ``True``, a :class:`numpy.ndarray` will the edge endpoints be
+        returned instead.
 
     Returns
     -------
-    djk_iterator : An iterator over the edges in Dijkstra order.
+    dfs_iterator : Iterator or :class:`numpy.ndarray`
+        An iterator over the edges in Dijkstra order. If ``array == True``,
+        this will be a :class:`numpy.ndarray` instead, of shape ``(E,2)``,
+        containing the edge endpoints.
 
     See Also
     --------
@@ -1022,6 +1048,7 @@ def dijkstra_iterator(g, weight, source=None, dist_map=None, combine=None,
         graphs", Numerische Mathematik, 1:269-271, 1959.
     .. [dijkstra-bgl] http://www.boost.org/doc/libs/release/libs/graph/doc/dijkstra_shortest_paths_no_color_map.html
     .. [dijkstra-wikipedia] http://en.wikipedia.org/wiki/Dijkstra's_algorithm
+
     """
 
     if dist_map is None:
@@ -1046,7 +1073,14 @@ def dijkstra_iterator(g, weight, source=None, dist_map=None, combine=None,
     else:
         source = int(source)
     if compare is None and combine is None:
-        return libgraph_tool_search.dijkstra_generator_fast(g._Graph__graph,
+        if not array:
+            return libgraph_tool_search.dijkstra_generator_fast(g._Graph__graph,
+                                                                source,
+                                                                _prop("v", g, dist_map),
+                                                                _prop("e", g, weight),
+                                                                zero, infinity)
+        else:
+            return libgraph_tool_search.dijkstra_array_fast(g._Graph__graph,
                                                             source,
                                                             _prop("v", g, dist_map),
                                                             _prop("e", g, weight),
@@ -1056,7 +1090,15 @@ def dijkstra_iterator(g, weight, source=None, dist_map=None, combine=None,
             compare = lambda a, b: a < b
         if combine is None:
             combine = lambda a, b: a + b
-        return libgraph_tool_search.dijkstra_generator(g._Graph__graph,
+        if not array:
+            return libgraph_tool_search.dijkstra_generator(g._Graph__graph,
+                                                           source,
+                                                           _prop("v", g, dist_map),
+                                                           _prop("e", g, weight),
+                                                           compare, combine,
+                                                           zero, infinity)