Dwango Programming Contest V - Program B. Sum AND Subarrays

beta.atcoder.jp

とりあえず遅いとわかりつつも総当たり

 from itertools import combinations
 def slow_solve(N, K, As):
     alls = []
     for i in range(N):
         for j in range(i + 1, N + 1):
             alls.append(sum(As[i:j]))
     m = 0
     for c in combinations(alls, K):
         a = 0xFFFFFFFFFFF
         for p in c:
             a &= p
         m = max(m, a)
     return m

で、ここから高速化を考えていく。

全体を通して、N(N+1)/2個に対する美しさの計算部分と K個選択して最大を取る部分の2か所に分かれる。

前半は各美しさに対してまとめられる要素がなさそうなのであきらめて、 後半の高速化を考える。ビット毎の論理積を計算したときに最大となるためには、 ANDを取った時になるべく左側にビットが立つようにしてあげればよい。 それより右がどうなろうとも、そこで最大値が決まる。

bit計算をすることによって1つの美しさに対して、O(log N)に計算量が落とせる。 全体としてはO(N2 * log N )

美しさの計算段階でまず2進数にする。 今回は109 * 1000 = 1012 = E8 D4A5 1000 < FF FFFF FFFF = 40bitだけあれば十分。 ここを109で計算していて時間をロスした。

ではここでどのようにK個を取るかを考える。

以下の4つをベースにKをずらしながら考える。

1111
1001
0111
0110

先ほどの通り、左のbitから一つずつ見ていく。

まず、K=3の時、一番左のbitを立てることは不可能であるので、 次のbitに行き、ビットが立っている行だけを取ればよい。

K=2の時、一番左のビットが立てられれば良いので、上の二つを取ればよい。

K=1の時、一番左のビットを見ている段階では、上の二つのうちどちらを取ればよいかが不明であるが、 上の二つのうちどちらかを取れば良いことがわかるのでこれらをターゲット、かつ次のbitから検索をかければよい。

まとめると再帰がいけそうなので、検索bit位置と、検索対象列を引数に 関数 find(i, alls) として、i番目のbitの立っている数と比較、K個ピッタリでbitが立っている場合を探す。

また、K個以上bitが立っていた場合、そのbitが立っている美しさに絞ったうえで、 次のbitに検索をかけるが、これ以降で一つも見つからない可能性がある。 それはすなわちどれをとっても同じ(次のbit以降はすべて0にしかなりえない)ということなので、適当に最初からK個を取る。

見つけたものはすでにMAXの保証があるので、同じようにAND計算してあげればよい。

def solve(N, K, As):
    alls = []
    for i in range(N):
        for j in range(i + 1, N + 1):
            alls.append("{:041b}".format(sum(As[i:j])))
    def find(i, alls):
        if i > 40:
            return []
        ones = [a for a in alls if a[i] == "1"]
        if len(ones) < K:
            return find(i + 1, alls)
        elif len(ones) == K:
            return ones
        elif len(ones) > K:
            ret = find(i + 1, ones)
            if len(ret) == 0:
                return ones[0:K]
            else:
                return ret
 
    ans = find(0, alls)
    if len(ans):
        a = 0xFFFFFFFFFFF
        for p in ans:
            a &= int(p, 2)
        return a
    else:
        return 0