-
Notifications
You must be signed in to change notification settings - Fork 0
/
bruteforce.py
130 lines (106 loc) · 4.31 KB
/
bruteforce.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
#!/usr/bin/env python3
"""
Brute force implementation
~~~~~~~~~~~~~~~~~~~~~~~~~~
Determines all possible arborescences in a given rooted digraph.
Returns the one of min-weight. Arborescences are found by
considering all edges extending the tree retrieved so far
until the tree size hits the total number of vertices.
Helps to verify testcases. Not meant for any other purposes.
(C) 2016, CC-0, Lukas Prokop
"""
import os
import sys
import argparse
import itertools
def read_input_graph(filepath: str):
"""Given a filepath, read a digraph file.
The format is specified in the script documentation.
:param filepath: A filepath of a digraph file
:type filepath: str
:return: set of vertices, edges and a root vertex
:rtype: ([int, ...], [(int, int, float), ...], int)
"""
vertices = []
edges = []
root = None
first = True
with open(filepath, encoding='utf-8') as fd:
for lineno, line in enumerate(fd):
if line.startswith("#"):
continue
if first:
vals = tuple(map(int, line.split()))
assert len(vals) >= 3, "first line must contain 3 integers"
assert vals[0] > 0, "number of vertices must be positive"
assert vals[1] >= 0, "number of edges must be non-negative"
assert vals[2] > 0, "root must be an existing vertex"
vertices = list(range(1, vals[0] + 1))
num_edges = vals[1]
root = vals[2]
first = False
else:
vals = line.split()
assert len(vals) == 3, "every edge line must contain 3 values"
assert int(vals[0]) > 0 and int(vals[1]) > 0, \
"vertices must be 1-enumerated (1..n)"
edges.append((int(vals[0]), int(vals[1]), float(vals[2])))
assert not first, "file must not be empty"
assert len(edges) == num_edges, "Actual # of edges differs from specified"
assert root in vertices, "root id exceeds vertex enumeration"
assert all(s in vertices and d in vertices for (s, d, w) in edges)
return (vertices, edges, root)
def find_arborescences(V, E, root, base=None):
"""Retrieve all arborescences.
:param V: set of vertices
:type V: {int, int, ...}
:param E: set of edges
:type E: {(int, int, float), ...}
:param base: an intermediate base tree retrieved
:type base: [(int, int, float), ...]
:return: generator for arborescences
:rtype: [[(int, int, float), ...]]
"""
if not base:
base = []
nodes = {e[0] for e in base}.union({e[1] for e in base})
if len(nodes) == len(V):
yield base
return
if not nodes:
nodes = [root]
for node in nodes:
for edge in filter(lambda e: e[0] == node, E):
if edge[1] in nodes:
# makes it cyclic
continue
yield from find_arborescences(V, E, root, base + [edge])
def main(filepath, print_base=False):
"""Main routine.
:param filepath: A filepath to a digraph file
:type filepath: str
:param print_base: Shall I print the base graph as 'b' lines?
:type print_base: bool
"""
V, E, root = read_input_graph(filepath)
max_branching = []
total_weight = float('inf')
for arb in find_arborescences(V, E, root):
weight = sum([e[2] for e in arb])
if weight < total_weight:
max_branching = arb
total_weight = weight
print('{} {} {} {}'.format(max(V), len(max_branching), root, total_weight))
if print_base:
for (s, d, w) in E:
print('b {} {} {}'.format(s, d, int(w) if w % 1 == 0.0 else w))
for (s, d, w) in max_branching:
print('{} {} {}'.format(s, d, int(w) if w % 1 == 0.0 else w))
if __name__ == '__main__':
d = 'Brute-force implementation to find min-weight arborescence.'
parser = argparse.ArgumentParser(description=d)
parser.add_argument('digraphfile', help='source file to read graph from')
parser.add_argument('-b', '--base', action='store_true',
help='print base graph as "b " lines')
args = parser.parse_args()
main(args.digraphfile, print_base=args.base)