Codeforces Round #612 (Div2) B Hyperset

年末年始休みだったのでcodeforcesに手を出してみましたがBでMLE,TLEを出し続け撃沈しました。pythonのような遅い言語には厳しいという噂は聞いていましたが本当のようです。（あとnumpy使えないのも提出してから知った...）

https://codeforces.com/contest/1287

SETT
TEST
EEET
ESTE
STES

のような文字列が与えられ、各列の文字の種類が1種か3種になるような3行の抜き出し方を数える問題。

行数N, 列数Mに対して制約はN<1.5*10**3, M<30, Time limit : 3s O(N*N*M) な2行固定の探索を考えましたが= 6 * 10**7 でpythonでは普通には間に合わないです。numpyもコドフォでは使えないことを提出後知りました。事前計算などを組み合わせたりbit演算で実現したりして高速化を図るも撃沈。

具体的に取ろうとした手法は以下です。S:0,E:1,T:2と数値化して各行をベクトルとしてみなすと、条件を満たす3行の和は%3でゼロベクトルとなっているはず。なので2行固定して、a+bを計算し、dictを使えば残りの1行はO(1)でその数がわかります。各行はdistinctなので残りの1行と固定した2行が重複してしまうことはないです。

from collections import defaultdict
N, M = map(int, input().split())
xs = []
d = defaultdict(int)
for _ in range(N):
 x = [(2 if c== "E" else 3) if c != "S" else 1 for c in input().strip()]
 xs.append(x)
 y = tuple(3-c for c in x)
 d[y] += 1
print(sum(d[tuple((a+b)%3 for a, b in zip(xs[i],xs[j]))] for i in range(N) for j in range(i+1, N))//3)

終了後、 S : 0010000, E : 00001000, T : 00010000 のようにbit化して連結したものを各行として利用して強引に上記アルゴリズムをbit演算に落としたら通りました。 多分めっちゃ不自然な解法なのですが、効率的にbit演算に落とすテクニックがわからないのでやむなし。

from collections import defaultdict
N, M = map(int, input().split())
xs = []
di = defaultdict(int)
for _ in range(N):
    x = 0
    for i, c in enumerate(input().strip()):
        x += ((4 if c== "E" else 1) if c != "S" else 2) << (7*i+2)
    xs.append(x)
    di[x] += 1
m= 0
for i in range(M):
    m += (7<<2)<<(7*i)
r = 0
for i, a in enumerate(xs):
    for j in range(i+1, N):
        b = xs[j]
        c = ((a&b)<<1)|((a&b)<<2)|((a&b)>>1)|((a&b)>>2)
        r += di[(((~(a^b))^c)&m)]
print(r//3)