Merge pull request #22 from LucienShui/feature/refact_and_add_edge

Remove extract graph, add Edge, refact code

flow_network/__init__.py

@@ -1,6 +1,6 @@
 from __future__ import absolute_import, print_function
 
-from flow_network.flow_network import MaximumFlow, MinimumCostFlow
+from .network import MaximumFlow, MinimumCostFlow
 
 from .__version__ import __title__, __description__, __url__
 from .__version__ import __version__, __build__, __author__

flow_network/__version__.py

@@ -1,7 +1,7 @@
 __title__ = 'flow-network'
 __description__ = 'Flow Network Python Library'
 __url__ = 'https://github.com/LucienShui/flow-network'
-__version_info__ = (0, 1, 10)
+__version_info__ = (0, 1, 11)
 __version__ = '.'.join(map(str, __version_info__))
 __build__ = eval(f"0x{''.join(map(lambda x: f'{int(x):02d}', __version_info__))}")
 __author__ = 'Lucien Shui'

flow_network/edges.py

@@ -0,0 +1,12 @@
+class Edge(object):
+    def __init__(self, u: int, v: int, capacity: int, flow: int = 0):
+        self.u: int = u
+        self.v: int = v
+        self.capacity: int = capacity
+        self.flow: int = flow
+
+
+class EdgeWithCost(Edge):
+    def __init__(self, u: int, v: int, capacity: int, cost: int, flow: int = 0):
+        super().__init__(u, v, capacity, flow)
+        self.cost: int = cost

flow_network/network.py

@@ -0,0 +1,123 @@
+from .core import CBaseNetwork, CMinimumCostFlow
+from .pycore import PyMinimumCostFlow
+from .edges import EdgeWithCost
+from .core import CMaximumFlow
+from .edges import Edge
+
+import typing
+
+
+class NetWork:
+
+    def __init__(self, n: int, algorithm_name: str, c_backend, python_backend=None, backend: str = 'c'):
+        super().__init__()
+
+        self._algorithm_name = algorithm_name
+        self._n: int = n
+
+        if backend.lower() == 'python' and python_backend is not None:
+            self._obj: CBaseNetwork = python_backend(n)
+        else:
+            self._obj: CBaseNetwork = c_backend(n)
+
+        self.edges: typing.List[Edge] = ...
+
+    def _run(self, s: int, t: int) -> [typing.Tuple, int]:
+        """
+        inference
+        :param s: source point's index
+        :param t: target point's index
+        :return: tuple or int
+        """
+        result = self._obj.run(s, t)
+
+        for i in range(0, len(self._obj.graph.edges), 2):
+            self.edges[i >> 1].flow = self._obj.graph.edges[i ^ 1].flow
+
+        return result
+
+    def _add_edge(self, u: int, v: int, *args) -> None:
+
+        if not 1 <= len(args) <= 2:
+            raise AssertionError('length of args must >= 1 and <= 2')
+
+        for arg in (u, v) + args:
+            if not isinstance(arg, int):
+                raise AssertionError(f'every arg should be type of int, got {type(arg)}')
+
+        for node in ['u', 'v']:
+            if eval(node) < 0:
+                raise AssertionError(f'index of {node} is {eval(node)}, which should be greater or equal to 0')
+
+            if eval(node) >= self._n:
+                raise AssertionError(f'index of {node} is {eval(node)}, which should be less than n = {self._n}')
+
+        self._obj.graph.add_edge(u, v, *args)
+
+    def summary(self, line_length: int = 32, print_fn=print):
+        vertex_cnt: int = self._n
+        edge_cnt: int = len(self.edges)
+
+        print_fn(f'''{'=' * line_length}
+{' '.join([each.capitalize() for each in self._algorithm_name.split('_')])}
+{'-' * line_length}
+Number of vertices: {vertex_cnt}
+Number of edges: {edge_cnt}
+{"=" * line_length}''')
+
+
+class MaximumFlow(NetWork):
+
+    def __init__(self, n: int, backend: str = 'c'):
+        super().__init__(n, 'flow_network', CMaximumFlow, None, backend)
+        self.edges: typing.List[Edge] = []
+        self._obj: CMaximumFlow = CMaximumFlow(n)
+
+    def add_edge(self, u: int, v: int, flow: int) -> None:
+        """
+        add edge from u to v with flow and cost
+        :param u: point's index
+        :param v: point's index
+        :param flow: edge capacity
+        :return: None
+        """
+        self._add_edge(u, v, flow)
+        self.edges.append(Edge(u, v, flow))
+
+    def run(self, s: int, t: int) -> int:
+        """
+        inference
+        :param s: source point's index
+        :param t: target point's index
+        :return: flow
+        """
+        return self._run(s, t)
+
+
+class MinimumCostFlow(NetWork):
+
+    def __init__(self, n: int, backend: str = 'c'):
+        super().__init__(n, 'minimum_cost_flow', CMinimumCostFlow, PyMinimumCostFlow, backend)
+        self.edges: typing.List[EdgeWithCost] = []
+
+    def add_edge(self, u: int, v: int, flow: int, cost: int) -> None:
+        """
+        add edge from u to v with flow and cost
+        :param u: point's index
+        :param v: point's index
+        :param flow: edge capacity
+        :param cost: cost for cutting an edge
+        :return: None
+        """
+        self._add_edge(u, v, flow, cost)
+        self.edges.append(EdgeWithCost(u, v, flow, cost))
+
+    def run(self, s: int, t: int) -> (int, int):
+        """
+        inference
+        :param s: source point's index
+        :param t: target point's index
+        :return: flow, cost
+        """
+        flow, cost = self._run(s, t)
+        return flow, cost

flow_network/pycore/__init__.py

@@ -0,0 +1 @@
+from .py_minimum_cost_flow import PyMinimumCostFlow

tests.py

@@ -25,8 +25,9 @@ def maximum_flow_test(self, backend: str = 'c'):
 
         self.assertEqual(1, maximum_flow.run(0, 4))
 
-        self.assertEqual(0, maximum_flow.edges[0][2])
-        self.assertEqual(1, maximum_flow.edges[1][2])
+        expected_flow = [1, 0, 0, 1, 1]
+
+        self.assertEqual(expected_flow, [edge.flow for edge in maximum_flow.edges])
 
     def minimum_cost_flow_test(self, backend: str = 'c'):
         minimum_cost_flow = MinimumCostFlow(5, backend)
@@ -44,27 +45,21 @@ def minimum_cost_flow_test(self, backend: str = 'c'):
             minimum_cost_flow.add_edge(u, v, flow, cost)
 
         minimum_cost_flow.summary()
-        minimum_cost_flow.extract_graph('graph.txt')
 
         flow, cost = minimum_cost_flow.run(0, 4)
 
         self.assertEqual(1, flow)
         self.assertEqual(4, cost)
 
-        self.assertEqual(0, minimum_cost_flow.edges[0][2])
-        self.assertEqual(1, minimum_cost_flow.edges[1][2])
+        expect_flow = [1, 1, 1, 0, 0]
+
+        self.assertEqual(expect_flow, [edge.flow for edge in minimum_cost_flow.edges])
 
     def test_flow_network(self):
         for backend in ['c', 'python']:
             self.maximum_flow_test(backend)
             self.minimum_cost_flow_test(backend)
 
-    def test_tuple_modifier(self):
-        from flow_network.util import tuple_modifier
-        buf = (1, 2, 3)
-        result = tuple_modifier(buf, 1, 4)
-        self.assertEqual((1, 4, 3), result)
-
 
 if __name__ == '__main__':
     unittest.main()