1import numpy as np
 2META = {'name':'corrupted_root_graph','domain':'graph-nn','description':'Root localization on permuted recursive trees with shortcut corruption.'}
 3def _one(rng,n=32,shortcuts=8):
 4    edges=[(v,int(rng.integers(v))) for v in range(1,n)]
 5    used={tuple(sorted(e)) for e in edges}
 6    for _ in range(shortcuts):
 7        for _ in range(100):
 8            a,b=map(int,rng.integers(n,size=2)); e=tuple(sorted((a,b)))
 9            if a!=b and e not in used: break
10        used.add(e); edges.append(e)
11    perm=rng.permutation(n); A=np.zeros((n,n),np.float32)
12    for a,b in edges: a,b=int(perm[a]),int(perm[b]); A[a,b]=A[b,a]=1
13    return A,int(np.where(perm==0)[0][0])
14def get_dataset(seed,n_train,n_test):
15    rng=np.random.default_rng(seed); n=32
16    def make(k):
17        x=np.empty((k,n,n),np.float32); y=np.empty(k,np.int64)
18        for i in range(k): x[i],y[i]=_one(rng,n)
19        return x,y
20    xtr,ytr=make(n_train); xte,yte=make(n_test)
21    return {'xtr':xtr,'ytr':ytr,'xte':xte,'yte':yte,'task':'classification','metric':'cross_entropy'}