main.py (1193B)
1 #!/usr/bin/python3 2 3 # compute undirected graph of configurations for efficient multinomial 4 # coefficient computation in SWERC'16 problem H 5 6 # use as: ./main.py MASS DIMENSIONS > graph.dot 7 # then compile graph with: dot -Tps graph.dot > graph.ps 8 9 import sys 10 11 def tples(n, h, l=None): 12 if l == None: 13 l = n+1 14 if n == 0: 15 return [()] 16 if h == 0: 17 return [] 18 lo = [] 19 for i in range(min(n, l)): 20 r = tples(n-i-1, h-1, i+1) 21 for t in r: 22 lo.append(tuple([i+1] + list(t))) 23 return lo 24 25 def adj(l1, l2): 26 if abs(len(l1) - len(l2)) > 1: 27 return False 28 if len(l2) > len(l1): 29 return adj(l2, l1) 30 if len(l1) > len(l2): 31 l2a = list(l2) + [0] 32 else: 33 l2a = list(l2) 34 s = sum(abs(x - y) for (x, y) in zip(l1, l2a)) 35 return s == 2 36 37 if __name__ == '__main__': 38 ml = tples(int(sys.argv[1]), int(sys.argv[2])) 39 print("graph G {") 40 for i in range(len(ml)): 41 print('n%d[label="%s"];' % (i, ml[i])) 42 for i in range(len(ml)): 43 for j in range(len(ml)): 44 if i < j: 45 if adj(ml[i], ml[j]): 46 print("n%d -- n%d;" % (i, j)) 47 print("}") 48