fork download
  1. from itertools import zip_longest, count, permutations
  2. import re
  3.  
  4. def uniqc(xs):
  5. seen = set((None,))
  6. return [''.join('' if c in seen else seen.add(c) or c for c in w)
  7. for w in xs]
  8.  
  9. def solve(q):
  10. def term(var):
  11. return '+'.join("a[{0}]*{1}".format(index[x], 10**i)
  12. for i, x in enumerate(reversed(var)))
  13.  
  14. variables = re.findall(r'\w+', q)
  15. steps = uniqc(zip_longest(*variables))
  16. index = dict(zip(''.join(steps), count()))
  17. limits = [sum(int(10**(e - 1)) for e in t)
  18. for t in [[x - i for x in [len(v) for v in variables]]
  19. for i in range(len(steps))]]
  20. qq = re.sub('^([^=]+)=+(.*)$', r'(\1)-(\2)', q)
  21. e = 'lambda a: abs(%s)' % re.sub(r'\w+', lambda m: term(m.group()), qq)
  22. distance = eval(e)
  23. distancex = lambda a: distance((a + (5,)*10)[:len(index)])
  24.  
  25. def advance(steps, limits, c, r):
  26. if not steps:
  27. if not distance(c):
  28. yield re.sub(r'\w', lambda m: str(c[index[m.group()]]), q)
  29. elif not steps[0]:
  30. yield from advance(steps[1:], limits[1:], c, r)
  31. else:
  32. xs = r if c else set(range(1,10))
  33. ps = (c + p for p in permutations(xs, len(steps[0])))
  34. for d, p in sorted((distancex(p), p) for p in ps):
  35. if d < limits[0]:
  36. yield from advance(steps[1:], limits[1:], p, r - set(p))
  37.  
  38. yield from advance(steps, limits, (), set(range(10)))
  39.  
  40. testcase = '''
  41. SEND + MORE = MONEY
  42. APPLE + GRAPE = CHERRY
  43. KYOTO + OSAKA = TOKYO
  44. AOMORI + SAPPORO = FUKUOKA
  45. KOCHI + OSAKA = NAGOYA
  46. OKAYAMA + WAKAYAMA = KANAZAWA
  47. FUKUOKA + NAGANO = YOKOHAMA
  48. KAMAKURA + MORIOKA = NAGASAKI
  49. KUMAMOTO + TOTTORI = WAKAYAMA
  50. OKAYAMA + WAKAYAMA = YAMAGATA
  51. TOTTORI + YAMAGATA = YOKOHAMA
  52. '''
  53. for case in testcase.strip().splitlines():
  54. print(next(solve(case)))
  55.  
Runtime error #stdin #stdout #stderr 0.05s 9568KB
stdin
Standard input is empty
stdout
Standard output is empty
stderr
  File "prog.py", line 30
    yield from advance(steps[1:], limits[1:], c, r)
             ^
SyntaxError: invalid syntax