Python实现 Preflow Push (Push-Relabel) 算法求最大流 并用Cython优化提速

原生python运行速度很慢，只要数据量大于500，求解就变得十分困难

五组测试数据

测试用例的第一行为图的节点数和边数，第二行为最大流算法的起始节点和中止节点，剩余所有行均为有向加权边，其中前两个数字代表边的两个端点，后一个数字代表边的权重。

·测试用例

原生代码实现

思路是保持解的最优性，寻找解的可行性。
为实现此目的，为每个点赋予一个高度值，其中源点高度在初始时设置为节点数目。
源点具有无限容量（盈余），每个点只能向比自己高度低的点推流，且每次推流都尽可能大。
如果自己有盈余却没有任何点可以推流，则增加自己的高度（Relabel）后再次尝试。
算法直到除了源点和宿点外，所有点的盈余都为0（找到可行解）为止。

#!/usr/bin/env python3
#preflow.py
#fumiama 20201225

from time import time
from heapq import heappop, heappush, heapify

class Node(object):
    def __init__(self):
        self.remain = 0
        self.height = 0
    def relabel(self): self.height += 1
    def push(self, flow): self.remain += flow
    def pour(self, flow): self.remain -= flow
    #def __repr__(self): return "盈余: "+str(self.remain)+" 高度: "+str(self.height)
class Graph(object):
    def __init__(self, nodeSize):
        self.nodeSize = nodeSize
        self.node2 = dict()
        self.nodes = [Node() for i in range(nodeSize+1)]
    def addEdge(self, a, b, w):
        if a not in self.node2.keys(): self.node2[a] = []
        heappush(self.node2[a], (b, w))
    def addWeight(self, a, b, w):
        #print("addWeight:", a, b, w)
        if a not in self.node2.keys(): self.addEdge(a, b, w)
        else:
            notChange = 1
            for node in self.node2[a]:
                if node[0] == b:
                    index = self.node2[a].index(node)
                    #print("index:", index, "node:", node)
                    previousWeight = self.node2[a][index][1]
                    newWeight = previousWeight + w
                    if newWeight: self.node2[a][index] = (b, newWeight)
                    else: self.delEdge(a, b)
                    notChange = 0
                    break
            if notChange: self.addEdge(a, b, w)
    def delEdge(self, a, b):
        for node in self.node2[a]:
            if node[0] == b: self.node2[a].remove(node)
        heapify(self.node2[a])
    def getWeight(self, a, b):
        if a in self.node2.keys():
            for node in self.node2[a]:
                if node[0] == b: return node[1]
        return -1
    def findMaxFlow(self, s, t):
        stack = []
        self.nodes[s].height = self.nodeSize
        while self.node2[s]:
            node = self.node2[s].pop()                  #删除正向边
            self.nodes[node[0]].remain = node[1]
            stack.append(node[0])
            self.addWeight(node[0], s, node[1])      #反向加边
        ##print("初始点:", stack)
        while stack:
            #print("stack:", stack)
            node = stack.pop()
            nodeObj = self.nodes[node]
            #print("弹出:", node, nodeObj)
            if nodeObj.remain and node != t:
                noPush = 1
                #print("图:", self.node2)
                for nextNode in self.node2[node]:
                    #print("(", node, ", ", nodeObj.remain, ")->", nextNode, sep='')
                    nextNodeObj = self.nodes[nextNode[0]]
                    if nodeObj.height > nextNodeObj.height:
                        noPush = 0
                        if nodeObj.remain <= nextNode[1]:
                            nextNodeObj.push(nodeObj.remain)
                            self.addWeight(node, nextNode[0], -nodeObj.remain)
                            #print("addWeight:", -nodeObj.remain)
                            self.addWeight(nextNode[0], node, nodeObj.remain)
                            #print("addBack:", nodeObj.remain)
                            nodeObj.remain = 0
                            #print("非饱和推送到:", nextNode, "该点", nextNodeObj)
                        else:
                            nextNodeObj.push(nextNode[1])
                            nodeObj.pour(nextNode[1])
                            self.delEdge(node, nextNode[0])
                            self.addWeight(nextNode[0], node, nextNode[1])
                            stack.append(node)
                            #print("饱和推送到:", nextNode, "该点", nextNodeObj)
                        if nextNode[0] != t and nextNode[0] != s and nextNodeObj.remain: stack.append(nextNode[0])
                if noPush and self.node2[node]:
                    nodeObj.relabel()
                    stack.append(node)
                    #print("relabel:", node, "to", nodeObj.height)
            #print('')
        return self.nodes[t].remain

if __name__ == '__main__':
    nodeSize, edgeSize = map(int, input().split())
    s, t = map(int, input().split())
    g = Graph(nodeSize)
    while edgeSize:
        a, b, w = map(int, input().split())
        g.addEdge(a, b, w)
        edgeSize -= 1
    time_start = time()
    print("最大流:", g.findMaxFlow(s, t))
    print("用时:", time() - time_start)

原生实现的样例测试

$ ./preflow.py < 测试用例1.txt
最大流: 13
用时: 0.00015997886657714844
$ ./preflow.py < 测试用例2.txt
最大流: 38
用时: 0.001840829849243164
$ ./preflow.py < 测试用例3.txt
最大流: 869
用时: 0.011828899383544922
$ ./preflow.py < 测试用例4.txt
最大流: 17639
用时: 207.51951098442078
$ ./preflow.py < 测试用例5.txt
(十分钟内未计算完成)

可以看到，随着点的数量增加，求解所需时间迅速呈指数形式增加。但是相比Edmonds-Karp算法，速度已经有了不小的提升。
造成运行缓慢的一个主要原因是python原生代码的缓慢，因此考虑使用cython重写算法代码以加快求解速度。cython会将我们的代码编译为以c语言实现的python库供我们调用。
为此，需要编写以下文件，并手动编译库文件：

pyx文件（代码主体）

#cython: language_level=3
#preflow_c.pyx
#fumiama 20201226

from time import time
from libc.string cimport memcpy, memset
from libc.stdlib cimport malloc, free

cdef int nodeSize
cdef int** arr

cdef set(int x, int y, int item):
    global arr
    #print("set", x, y, item)
    arr[x][y] = item
cdef int get(int x, int y):
    global arr
    #print("get", x, y)
    return arr[x][y]
cpdef createArray(int x, int y):
    global arr
    cdef int i
    arr = <int**>malloc(x * sizeof(int*))
    for i in range(x):
        arr[i] = <int*>malloc(y * sizeof(int))
        memset(arr[i], 0, y * sizeof(int))

cdef int* stack
cdef int sp = 0
cdef int* touch
cdef initStack(int length):
    global stack, sp, touch
    stack = <int*>malloc(length)
    touch = <int*>malloc(length)
    memset(touch, 0, length)
    sp = 0
cdef push(int item):
    global stack, sp, touch
    stack[sp] = item
    touch[item] = 1
    sp += 1
cdef int pop():
    global stack, sp, touch
    cdef int re
    sp -= 1
    re = stack[sp]
    touch[re] = 0
    return re
cdef int inStack(int item):
    global touch
    return touch[item]
#cdef printStack():
#    global stack, sp
#    cdef int i
#    print("stack:[", end=' ')
#    for i in range(sp): print(stack[i], end=" ")
#    print(']')

cdef int* remain
cdef int* height
cdef initGraph(int nSize):
    global nodeSize, arr, remain, height
    nodeSize = nSize
    length = nSize * sizeof(int)
    createArray(nSize, nSize)
    remain = <int*>malloc(length)
    height = <int*>malloc(length)
    memset(remain, 0, length)
    memset(height, 0, length)
cdef addEdge(int a, int b, int w): set(a-1, b-1, w)
cdef addWeight(int a, int b, int w): set(a, b, w + get(a, b))
cdef delEdge(int a, int b): set(a, b, 0)

cdef int findMaxFlow(int s, int t):
    global nodeSize, height, sp
    cdef int length = nodeSize * sizeof(int)
    cdef int sM1 = s-1
    cdef int tmp, noPush, nextNode, notEmpty
    initStack(length)
    #print("Stack initialized.")
    height[sM1] = nodeSize
    for nextNode in range(nodeSize):
        tmp = get(sM1, nextNode)
        if tmp:
            #print("设置出边:", s, nextNode+1, tmp)
            set(sM1, nextNode, 0)
            remain[nextNode] = tmp
            push(nextNode)
            addWeight(nextNode, sM1, tmp)
    #print("Pushed from souce")
    while sp:
        #printStack()
        node = pop()
        #print(sp, end=' ')
        if remain[node] and node != t-1:
            noPush = 1
            for nextNode in range(nodeSize):
                tmp = get(node, nextNode)
                #print("inStack", nextNode, inStack(nextNode))
                if tmp and height[node] > height[nextNode]:
                    noPush = 0
                    #print("找到边:", node, nextNode, tmp)
                    if remain[node] <= tmp:
                        if nextNode != s-1: remain[nextNode] += remain[node]        #nextNodeObj.push(nodeObj.remain)
                        addWeight(node, nextNode, -remain[node])
                        addWeight(nextNode, node, remain[node])
                        remain[node] = 0
                    else:
                        if nextNode != s-1: remain[nextNode] += tmp
                        remain[node] -= tmp
                        delEdge(node, nextNode)
                        addWeight(nextNode, node, tmp)
                        if not inStack(node) and node != t-1: push(node)
                    #print(remain[node], end=' ')
                    if nextNode+1 != t and nextNode+1 != s and remain[nextNode] and not inStack(nextNode): push(nextNode)
            notEmpty = 0
            for nextNode in range(nodeSize):
                tmp = get(node, nextNode)
                if tmp and nextNode+1 != s:
                    notEmpty = 1
                    break
            if noPush and notEmpty and remain[node]:
                height[node] += 1
                push(node)
    return remain[t-1]

def main():
    cdef int nodeSize, edgeSize, s, t, a, b, w
    cdef double time_start
    nodeSize, edgeSize = map(int, input().split())
    s, t = map(int, input().split())
    initGraph(nodeSize)
    #print("Graph initialized.")
    while edgeSize:
        a, b, w = map(int, input().split())
        #print("Add edge:", a, b, w)
        addEdge(a, b, w)
        edgeSize -= 1
    time_start = time()
    print("最大流:", findMaxFlow(s, t))
    print("用时:", time() - time_start)

setup文件（构建库）

# cython: language_level=3
from distutils.core import setup
from Cython.Build import cythonize

setup(
    ext_modules = cythonize("preflow_c.pyx")
)

编译

$ python3 ./preflow_c_setup.py build_ext --inplace

调用库

#edmonds_c_run.py
import preflow_c
preflow_c.main()

进行测试

$ ./preflow_c_run.py < 测试用例1.txt
最大流: 13
用时: 3.0994415283203125e-05
$ ./preflow_c_run.py < 测试用例2.txt
最大流: 38
用时: 6.29425048828125e-05
$ ./preflow_c_run.py < 测试用例3.txt
最大流: 869
用时: 0.0002541542053222656
$ ./preflow_c_run.py < 测试用例4.txt
最大流: 17639
用时: 0.7053031921386719
$ ./preflow_c_run.py < 测试用例5.txt
最大流: 62401
用时: 4.692018032073975

可见在使用cython优化之后，运行速度有了极大提升，基本可以胜任大部分场景下的计算需求。