AtCoder Beginner Contest 247 E,G,Ex問題メモ

from itertools import groupby

n, x, y = map(int, input().split())
aaa = list(map(int, input().split()))

if x == y:
    ans = 0
    for a, itr in groupby(aaa):
        if a == x:
            l = len(list(itr))
            ans += l * (l + 1) // 2
    print(ans)
    exit()

x_idx = [-1]
y_idx_l = []
z_idx_l = []
last_y_idx = -1
last_z_idx = -1
for i, a in enumerate(aaa):
    if a == x:
        x_idx.append(i)
        y_idx_l.append(last_y_idx)
        z_idx_l.append(last_z_idx)
    elif a == y:
        last_y_idx = i
    elif a > x or a < y:
        last_z_idx = i
y_idx_r = []
z_idx_r = []
last_y_idx = n
last_z_idx = n
for i, a in reversed(list(enumerate(aaa))):
    if a == x:
        y_idx_r.append(last_y_idx)
        z_idx_r.append(last_z_idx)
    elif a == y:
        last_y_idx = i
    elif a > x or a < y:
        last_z_idx = i
y_idx_r.reverse()
z_idx_r.reverse()

ans = 0
for i, xi in enumerate(x_idx[1:]):
    lxi = x_idx[i]
    lyi = y_idx_l[i]
    ryi = y_idx_r[i]
    lzi = z_idx_l[i]
    rzi = z_idx_r[i]

    l_pat1 = max(0, lyi - max(lxi, lzi))
    r_pat1 = rzi - xi
    l_pat2 = max(0, xi - max(lxi, lyi, lzi))
    r_pat2 = max(0, rzi - ryi)
    ans += l_pat1 * r_pat1 + l_pat2 * r_pat2

print(ans)

手元に持ってる最小費用流のライブラリがDijkstraでなくBellman-Ford実装にしてたため、TLEが取り切れなかった。

ちゃんとDijkstra版のライブラリも持っときましょうね。（ちゃんと理解してたらその場で書けるだろって話ではある）

numpy.fftを使った畳み込みではTLEが取り切れなかった。

（numpy.fftは小数で演算されるため、オーバーフローする可能性のある畳み込みでは、bitをずらして3回行う必要がある）

まぁ、それ以外にもよろしくない実装があったのかも知れないが、畳み込みの処理を直書きするのはやはり重たい。

import os
import sys
from heapq import heapify, heappop, heapreplace

import numpy as np


def solve(inp):
    def bit_length(n):
        x = 0
        while 1 << x < n:
            x += 1
        return x

    def bit_scan_forward(n):
        x = 0
        while n & 1 == 0:
            n >>= 1
            x += 1
        return x

    def pow_mod(x, n, m):
        r = 1
        y = x % m
        while n:
            if n & 1:
                r = (r * y) % m
            y = y * y % m
            n >>= 1
        return r

    def get_primitive_root(m):
        if m == 2: return 1
        if m == 167772161: return 3
        if m == 469762049: return 3
        if m == 754974721: return 11
        if m == 998244353: return 3
        divs = [2]
        x = (m - 1) // 2
        while x & 1 == 0:
            x >>= 1
        i = 3
        while i * i <= x:
            if x % i == 0:
                divs.append(i)
                x //= i
                while x % i == 0:
                    x //= i
            i += 2
        if x > 1:
            divs.append(x)
        g = 2
        while True:
            ok = True
            for d in divs:
                if pow_mod(g, (m - 1) // d, m) == 1:
                    ok = False
                    break
            if ok:
                return g

    def butterfly_prepare(mod, primitive_root):
        sum_e = np.zeros(30, np.int64)
        sum_ie = np.zeros(30, np.int64)
        es = np.zeros(30, np.int64)
        ies = np.zeros(30, np.int64)
        cnt2 = bit_scan_forward(mod - 1)
        e = pow_mod(primitive_root, (mod - 1) >> cnt2, mod)
        ie = pow_mod(e, mod - 2, mod)
        for i in range(cnt2, 1, -1):
            es[i - 2] = e
            ies[i - 2] = ie
            e = e * e % mod
            ie = ie * ie % mod
        now_e = 1
        now_ie = 1
        for i in range(cnt2 - 1):
            sum_e[i] = es[i] * now_e % mod
            sum_ie[i] = ies[i] * now_ie % mod
            now_e = now_e * ies[i] % mod
            now_ie = now_ie * es[i] % mod
        return sum_e, sum_ie

    def butterfly(aaa, mod, sum_e):
        n = aaa.size
        h = bit_length(n)

        for ph in range(1, h + 1):
            w = 1 << (ph - 1)
            p = 1 << (h - ph)
            now = 1
            for s in range(w):
                offset = s << (h - ph + 1)
                for i in range(p):
                    l = aaa[i + offset]
                    r = aaa[i + offset + p] * now % mod
                    aaa[i + offset] = (l + r) % mod
                    aaa[i + offset + p] = (l - r) % mod
                now = now * sum_e[bit_scan_forward(~s)] % mod

    def butterfly_inv(aaa, mod, sum_ie):
        n = aaa.size
        h = bit_length(n)

        for ph in range(h, 0, -1):
            w = 1 << (ph - 1)
            p = 1 << (h - ph)
            inow = 1
            for s in range(w):
                offset = s << (h - ph + 1)
                for i in range(p):
                    l = aaa[i + offset]
                    r = aaa[i + offset + p]
                    aaa[i + offset] = (l + r) % mod
                    aaa[i + offset + p] = ((l - r) * inow) % mod
                inow = inow * sum_ie[bit_scan_forward(~s)] % mod

    def convolve_ntt(aaa, bbb, MOD):
        n = aaa.size
        m = bbb.size
        k = n + m - 1
        z = 1 << bit_length(k)
        raaa = np.zeros(z, np.int64)
        rbbb = np.zeros(z, np.int64)
        raaa[:n] = aaa
        rbbb[:m] = bbb
        primitive_root = get_primitive_root(MOD)
        sum_e, sum_ie = butterfly_prepare(MOD, primitive_root)
        butterfly(raaa, MOD, sum_e)
        butterfly(rbbb, MOD, sum_e)
        for i in range(z):
            raaa[i] = raaa[i] * rbbb[i] % MOD
        butterfly_inv(raaa, MOD, sum_ie)
        iz = pow_mod(z, MOD - 2, MOD)
        for i in range(k):
            raaa[i] = raaa[i] * iz % MOD
        return raaa[:k]

    MOD = 998244353

    k = inp[1]
    ccc = inp[2:]

    counter = {}
    for c in ccc:
        if c in counter:
            counter[c] += 1
        else:
            counter[c] = 1

    targets = []
    hq = []

    for cnt in counter.values():
        for i in range(1, cnt):
            targets.append(np.array([1, i], dtype=np.int64))
            hq.append((2, len(targets) - 1))

    if len(hq) == 0:
        if k % 2 == 0:
            return 1
        else:
            return 0

    heapify(hq)

    while len(hq) > 1:
        c1, i1 = heappop(hq)
        c2, i2 = hq[0]
        targets[i2] = convolve_ntt(targets[i1], targets[i2], MOD)
        heapreplace(hq, (c1 + c2 - 1, i2))

    _, i = hq[0]
    pat = targets[i]

    if k % 2 == 0:
        ans = pat[0:k + 1:2].sum() % MOD
    else:
        ans = pat[1:k + 1:2].sum() % MOD
    return ans


SIGNATURE = '(i8[:],)'
if sys.argv[-1] == 'ONLINE_JUDGE':
    from numba.pycc import CC

    cc = CC('my_module')
    cc.export('solve', SIGNATURE)(solve)
    cc.compile()
    exit()

if os.name == 'posix':
    # noinspection PyUnresolvedReferences
    from my_module import solve
else:
    from numba import njit

    solve = njit(SIGNATURE, cache=True)(solve)
    print('compiled', file=sys.stderr)

inp = np.fromstring(sys.stdin.read(), dtype=np.int64, sep=' ')
ans = solve(inp)
print(ans)