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

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

五组测试数据

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

·测试用例1（将txt压缩为xz后使用base16384转为文本）

赍薥潨一伹筫弘倀噀忀一丐臖嫎七倀栗币瀰丢礁凹僺藘熐珃箖孨舉詹瑷俶嘴奊懫羾姨羖盾眅芽贁蒎膝悢艵戧婝榬繣哧譹侽烿犎滃幗徤俊俿圐嫌綉殾穫嫻带潶宜疹讥及煷洟側廫盃渍制谳珲澼瘀侄买縐丁誀妇礜徟褂一一丑杚

·测试用例2（将txt压缩为xz后使用base16384转为文本）

赍薥潨一伹筫弘倀噀忀一丐臖嫎乃攅泗币砰丢荲湽貅桠缐剧绚嵵囉屿篶折蟤櫆尘蓇憨总昮瘓杪廍蓌冇偔窵蚣糌安挅憅婞吸痊嘳朷梨傉劈臹癭蟲汦彆匓裰嵌捚讐續抒覈极俭愰哽楔荱訠早瀤媩搿胹烷呲悹桜秾秵嶜凊孭乳厢灰眠暵缫痑攠搨茌矓竦诸侼氛堽暨禭圌筺禐瀔垚絹愙俦蟇热究葭狆弢姏簾窽瀀媃豷耳姸臁硵脰瞉炦畦帑壄藛灻猃皁樟艑蜦暸莠稼殔冸盿篷荞晪熩詥楷晘瞀聜絟务竉嘬贔勊瞵澃婞詈戼兠睆腝岪擬秱崉貁忊種椄容胖稢珫漡圏梻豃嚓脻楉櫽儒漇篠蛉狼粚抐俎敚撑曀祻潱艌譗嚘吗凿汒爣神簪嶶瘗狸夳索碵薹皭褒尼焊秛四擤晟筻滞聑忖湈笂庀噉煪彀覊優可縖涥梐橫累籋莧慛架謨曖俻们晶袩积啀烥箻禎擢獣砪卒譂槨椈础羑珈拠屩瑳牁懁忺眳蟲寷備嚐聾榰嚞稽斮瘖禊忌胿簈赧蟱稚綎婓嚟禾囯害孜灆仁洿爬涍瞇旞暠蓊危脌塟蟰藬嵃菐乷臧簻漂毷桴亼蔂粻睥撃攖櫹絢划冺匹簪燴啷完層嶳亱瓸蕾硳嘬紵帑焥佘膲摤贜獱熧怸坔忎洞憂篲絜嵆橚噕璌噖謌卻爮斈菛嗦簈竧休萨汰箻岴熌虯梏衕应簊翞坄孀暣櫳俓楶攏嬶悷儃縠刯灋埍燕蔍曷害捩萛傈挅篻膨冘晳蝍罺蓁竿扎赍碂狩洉幋喋咁螃呩縡灦榀捀塼責湻憩忝璷捧狕刏繡棤悄慊笹挖睅塠潩台羜菞玘星短祦犽娸貔梉瀦呎徣偿燵暶裤娇剓啇們浣蔯挩乡痳菆厱尝蔛浤苚挳謙矯瞏尼缲猨抺係礒藪奠譙欀茆觬耥碈擛燰毦妽櫌帄谌礷档倁厤籸跶跜襚乒浄獧獛杠嘺濝烦墧豰賍壽毇姹琼睧諩灣裒椆芠筵簼癊惕倳吁貚佚膩裲倸畝峍姃荫薀薾摽痦暯撛妝硁茿耊妝圔垳垚蕥洯烜壧姞够琡吐挣色繿堂篧檞撨嶍瀂狽譾賍蠴璋枣帙嫔欇枵塗苟羿裎姕擭慭嬦莊搢謼傮莑緂硖瘙拼埉烑彄趼买挿枠娉佼搽衹匱梸粭滬媣熛堢皖噭肜觖僆羷曃窥蕎屙仑惐薗报岈碑焆楸僂娡奻溮壜秝篕潥娤圯嶎烘硦瀺傊倵膾矇蕧簇芋囯聩敎訾毼訔扴襐缹娛裨筋纩崶畹域忷僂語諻汎毨藯綊俞忏芦绣砽嚀燿曀挿徦倛舆狦橊灧岌晚潖刜劫洀猞焆憡緖磂垺秗根厁臾纲揤梎滝蟐晙悾櫎奟秞繯涀朅庆厙裟哚懆嚯巈櫚痽臲葘徳潈蘑稐廥罕欦愷囜祽癮楈焝櫾崪羮竐杈詿帇氾虂槓娠燬佋怑磅秐帨臑莛烽氫泘冹袋园普檼蕁藍蟯梠竲朂穣俩臯嗢斆纹唺波暕虨賵缁儝濋绱厑菷矆攽兲巈搆洀丣貔營慇緆一呜姘噀三柺坮穱呿稈一一久獨㴽㴽㴽

·测试用例3、4、5

原生代码实现

思路是保持解的最优性，寻找解的可行性。
为实现此目的，为每个点赋予一个高度值，其中源点高度在初始时设置为节点数目。
源点具有无限容量（盈余），每个点只能向比自己高度低的点推流，且每次推流都尽可能大。
如果自己有盈余却没有任何点可以推流，则增加自己的高度（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优化之后，运行速度有了极大提升，基本可以胜任大部分场景下的计算需求。