trie2dot.py (1488B)
1 #!/usr/bin/env python3 2 3 """Takes json as input with labels [value1, value2] and produces dot, 4 usage: trie2dot.py prefix value1 value2""" 5 6 import json 7 import sys 8 from math import log 9 10 trie = json.load(sys.stdin) 11 12 free_id = 0 13 14 def cget(d, k): 15 try: 16 if k in d.keys(): 17 return d[k] 18 else: 19 return 0 20 except AttributeError: 21 # we have a list, not a dictionary 22 # this happens after majoritytrie.py 23 if k in d: 24 return 1 25 else: 26 return 0 27 28 def int2strbyte(i): 29 s = hex(i).split('x')[1] 30 if len(s) == 1: 31 return '0' + s 32 else: 33 return s 34 35 def fraction2rgb(fraction): 36 n = int(255*fraction) 37 return int2strbyte(n)+'00'+int2strbyte(255 - n) 38 39 def total(x): 40 key, node = x 41 try: 42 return sum(node[0].values()) 43 except AttributeError: 44 # we have only one value, not a dictionary 45 return 1 46 47 def to_dot(trie, prefix=''): 48 global free_id 49 50 values, children = trie 51 my_id = free_id 52 free_id += 1 53 count = cget(values, v1) + cget(values, v2) 54 fraction = cget(values, v2) / count 55 56 print("%d [label=\"%s\",color=\"#%s\",penwidth=%d]" % (my_id, prefix, 57 fraction2rgb(fraction), 1+int(log(count)))) 58 59 for (key, child) in sorted(children.items(), key=total, reverse=True): 60 i = to_dot(child, prefix+key) 61 print("%d -> %d [penwidth=%d]" % (my_id, i, 62 1+int(log(total((None, child)))))) 63 64 return my_id 65 66 print("digraph G {\naspect=\"1\"\n") 67 prefix = sys.argv[1] 68 v1 = sys.argv[2] 69 v2 = sys.argv[3] 70 to_dot(trie, prefix) 71 print("}")