Red-White Fence (Codeforces 1251F)
Tags: Others, fft

Polycarp wants to build a fence near his house. He has $$$n$$$ white boards and $$$k$$$ red boards he can use to build it. Each board is characterised by its length, which is an integer.A good fence should consist of exactly one red board and several (possibly zero) white boards. The red board should be the longest one in the fence (every white board used in the fence should be strictly shorter), and the sequence of lengths of boards should be ascending before the red board and descending after it. Formally, if $$$m$$$ boards are used, and their lengths are $$$l_1$$$, $$$l_2$$$, ..., $$$l_m$$$ in the order they are placed in the fence, from left to right (let's call this array $$$[l_1, l_2, \dots, l_m]$$$ the array of lengths), the following conditions should hold:  there should be exactly one red board in the fence (let its index be $$$j$$$);  for every $$$i \in [1, j - 1]$$$ $$$l_i < l_{i + 1}$$$;  for every $$$i \in [j, m - 1]$$$ $$$l_i > l_{i + 1}$$$. When Polycarp will build his fence, he will place all boards from left to right on the same height of $$$0$$$, without any gaps, so these boards compose a polygon:  Example: a fence with $$$[3, 5, 4, 2, 1]$$$ as the array of lengths. The second board is red. The perimeter of the fence is $$$20$$$. Polycarp is interested in fences of some special perimeters. He has $$$q$$$ even integers he really likes (these integers are $$$Q_1$$$, $$$Q_2$$$, ..., $$$Q_q$$$), and for every such integer $$$Q_i$$$, he wants to calculate the number of different fences with perimeter $$$Q_i$$$ he can build (two fences are considered different if their arrays of lengths are different). Can you help him calculate these values?

Input: The first line contains two integers $$$n$$$ and $$$k$$$ ($$$1 \le n \le 3 \cdot 10^5$$$, $$$1 \le k \le 5$$$) — the number of white and red boards Polycarp has.The second line contains $$$n$$$ integers $$$a_1$$$, $$$a_2$$$, ..., $$$a_n$$$ ($$$1 \le a_i \le 3 \cdot 10^5$$$) — the lengths of white boards Polycarp has.The third line contains $$$k$$$ integers $$$b_1$$$, $$$b_2$$$, ..., $$$b_k$$$ ($$$1 \le b_i \le 3 \cdot 10^5$$$) — the lengths of red boards Polycarp has. All $$$b_i$$$ are distinct.The fourth line contains one integer $$$q$$$ ($$$1 \le q \le 3 \cdot 10^5$$$) — the number of special integers.The fifth line contains $$$q$$$ integers $$$Q_1$$$, $$$Q_2$$$, ..., $$$Q_q$$$ ($$$4 \le Q_i \le 12 \cdot 10^5$$$, every $$$Q_i$$$ is even) — the special integers Polycarp likes.

Output: For each $$$Q_i$$$, print one integer — the number of good fences with perimeter $$$Q_i$$$ Polycarp can build, taken modulo $$$998244353$$$.

Note: Possible fences in the first example denoted by their arrays of lengths (the length of the red board is highlighted):  with perimeter $$$6$$$: $$$[\textbf{2}]$$$;  with perimeter $$$8$$$: $$$[1, \textbf{2}]$$$, $$$[\textbf{2}, 1]$$$;  with perimeter $$$10$$$: $$$[1, \textbf{2}, 1]$$$, $$$[\textbf{4}]$$$;  with perimeter $$$12$$$: $$$[1, \textbf{4}]$$$, $$$[3, \textbf{4}]$$$, $$$[\textbf{4}, 1]$$$, $$$[\textbf{4}, 3]$$$;  with perimeter $$$14$$$: $$$[1, \textbf{4}, 1]$$$, $$$[1, \textbf{4}, 3]$$$, $$$[3, \textbf{4}, 1]$$$, $$$[3, \textbf{4}, 3]$$$, $$$[1, 3, \textbf{4}]$$$, $$$[\textbf{4}, 3, 1]$$$;  with perimeter $$$16$$$: $$$[1, \textbf{4}, 3, 1]$$$, $$$[3, \textbf{4}, 3, 1]$$$, $$$[1, 3, \textbf{4}, 1]$$$, $$$[1, 3, \textbf{4}, 3]$$$;  with perimeter $$$18$$$: $$$[1, 3, \textbf{4}, 3, 1]$$$. 

Sample Inputs:
5 2
3 3 1 1 1
2 4
7
6 8 10 12 14 16 18

5 5
1 2 3 4 5
1 2 3 4 5
4
4 8 10 14

Sample Outputs:
1
2
2
4
6
4
1

1
3
5
20

Solution (Python):
(standard_input, packages, output_together) = (1, 1, 1)
(dfs, hashing, read_from_file) = (0, 0, 0)
de = 1
if 1:
    if standard_input:
        import io, os, sys
        input = lambda : sys.stdin.readline().strip()
        import math
        inf = math.inf

        def I():
            return input()

        def II():
            return int(input())

        def MII():
            return map(int, input().split())

        def LI():
            return list(input().split())

        def LII():
            return list(map(int, input().split()))

        def LFI():
            return list(map(float, input().split()))

        def GMI():
            return map(lambda x: int(x) - 1, input().split())

        def LGMI():
            return list(map(lambda x: int(x) - 1, input().split()))
    if packages:
        from io import BytesIO, IOBase
        import random
        import os
        import bisect
        import typing
        from collections import Counter, defaultdict, deque
        from copy import deepcopy
        from functools import cmp_to_key, lru_cache, reduce
        from heapq import merge, heapify, heappop, heappush, heappushpop, nlargest, nsmallest
        from itertools import accumulate, combinations, permutations, count, product
        from operator import add, iand, ior, itemgetter, mul, xor
        from string import ascii_lowercase, ascii_uppercase, ascii_letters
        from typing import *
        BUFSIZE = 4096
    if output_together:

        class FastIO(IOBase):
            newlines = 0

            def __init__(self, file):
                self._fd = file.fileno()
                self.buffer = BytesIO()
                self.writable = 'x' in file.mode or 'r' not in file.mode
                self.write = self.buffer.write if self.writable else None

            def read(self):
                while True:
                    b = os.read(self._fd, max(os.fstat(self._fd).st_size, BUFSIZE))
                    if not b:
                        break
                    ptr = self.buffer.tell()
                    (self.buffer.seek(0, 2), self.buffer.write(b), self.buffer.seek(ptr))
                self.newlines = 0
                return self.buffer.read()

            def readline(self):
                while self.newlines == 0:
                    b = os.read(self._fd, max(os.fstat(self._fd).st_size, BUFSIZE))
                    self.newlines = b.count(b'\n') + (not b)
                    ptr = self.buffer.tell()
                    (self.buffer.seek(0, 2), self.buffer.write(b), self.buffer.seek(ptr))
                self.newlines -= 1
                return self.buffer.readline()

            def flush(self):
                if self.writable:
                    os.write(self._fd, self.buffer.getvalue())
                    (self.buffer.truncate(0), self.buffer.seek(0))

        class IOWrapper(IOBase):

            def __init__(self, file):
                self.buffer = FastIO(file)
                self.flush = self.buffer.flush
                self.writable = self.buffer.writable
                self.write = lambda s: self.buffer.write(s.encode('ascii'))
                self.read = lambda : self.buffer.read().decode('ascii')
                self.readline = lambda : self.buffer.readline().decode('ascii')
        sys.stdout = IOWrapper(sys.stdout)
    if dfs:
        from types import GeneratorType

        def bootstrap(f, stack=[]):

            def wrappedfunc(*args, **kwargs):
                if stack:
                    return f(*args, **kwargs)
                else:
                    to = f(*args, **kwargs)
                    while True:
                        if type(to) is GeneratorType:
                            stack.append(to)
                            to = next(to)
                        else:
                            stack.pop()
                            if not stack:
                                break
                            to = stack[-1].send(to)
                    return to
            return wrappedfunc
    if hashing:
        RANDOM = random.getrandbits(20)

        class Wrapper(int):

            def __init__(self, x):
                int.__init__(x)

            def __hash__(self):
                return super(Wrapper, self).__hash__() ^ RANDOM
    if read_from_file:
        file = open('input.txt', 'r').readline().strip()[1:-1]
        fin = open(file, 'r')
        input = lambda : fin.readline().strip()
        output_file = open('output.txt', 'w')

        def fprint(*args, **kwargs):
            print(*args, **kwargs, file=output_file)
    if de:

        def debug(*args, **kwargs):
            print('\x1b[92m', end='')
            print(*args, **kwargs)
            print('\x1b[0m', end='')

class Factorial:

    def __init__(self, N, mod) -> None:
        N += 1
        self.mod = mod
        self.f = [1 for _ in range(N)]
        self.g = [1 for _ in range(N)]
        for i in range(1, N):
            self.f[i] = self.f[i - 1] * i % self.mod
        self.g[-1] = pow(self.f[-1], mod - 2, mod)
        for i in range(N - 2, -1, -1):
            self.g[i] = self.g[i + 1] * (i + 1) % self.mod

    def fac(self, n):
        return self.f[n]

    def fac_inv(self, n):
        return self.g[n]

    def combi(self, n, m):
        if n < m or m < 0 or n < 0:
            return 0
        return self.f[n] * self.g[m] % self.mod * self.g[n - m] % self.mod

    def permu(self, n, m):
        if n < m or m < 0 or n < 0:
            return 0
        return self.f[n] * self.g[n - m] % self.mod

    def catalan(self, n):
        return (self.combi(2 * n, n) - self.combi(2 * n, n - 1)) % self.mod

    def inv(self, n):
        return self.f[n - 1] * self.g[n] % self.mod
mod = 998244353
fact = Factorial(10 ** 6, mod)

def main():
    (n, k) = MII()
    whites = LII()
    reds = LII()
    q = II()
    queries = LII()
    cnt = [0] * (3 * 10 ** 5 + 1)
    for x in whites:
        cnt[x] += 1
        if cnt[x] > 2:
            cnt[x] = 2
    ans = [0] * q
    for r in reds:
        (c1, c2) = (0, 0)
        for i in range(r):
            if cnt[i] == 1:
                c1 += 1
            elif cnt[i] == 2:
                c2 += 1
        c2 *= 2
        dp = [0] * (c1 + c2 + 1)
        dp[0] = 1
        if c1 or c2:
            dp[1] = 2 * c1 + c2
        for i in range(2, c1 + c2 + 1):
            dp[i] = fact.inv(i) * ((2 * c1 + c2 - 3 * (i - 1)) * dp[i - 1] + (2 * c1 + 2 * c2 - 2 * (i - 2)) * dp[i - 2]) % mod
        for i in range(q):
            if queries[i] > 2 * r and queries[i] // 2 - r - 1 <= c1 + c2:
                ans[i] += dp[queries[i] // 2 - r - 1]
                if ans[i] > mod:
                    ans[i] -= mod
    print('\n'.join((str(x) for x in ans)))
    return
t = 1
for _ in range(t):
    main()

Solution (Python):
import sys, os, io
input = io.BytesIO(os.read(0, os.fstat(0).st_size)).readline
mod = 998244353
r = [pow(3, mod >> i, mod) for i in range(23)]
ir = [pow(i, mod - 2, mod) for i in r]
u = [0] * (1 << 23)

def convolve(f, g):
    (lf, lg) = (len(f), len(g))
    c = lf + lg - 1
    x = c.bit_length()
    l = 1 << x
    for i in range(lf):
        u[i] = f[i]
    for i in range(lf, l):
        u[i] = 0
    for i in range(lg):
        u[i ^ l] = g[i]
    for i in range(lg, l):
        u[i ^ l] = 0
    m = l >> 1
    for k in range(x, 0, -1):
        r0 = r[k]
        for i in range(0, l, m << 1):
            y = 1
            for j in range(i, m + i):
                (u[j], u[j + m]) = ((u[j] + u[j + m]) % mod, (u[j] - u[j + m]) * y % mod)
                (u[j + l], u[j + m + l]) = ((u[j + l] + u[j + m + l]) % mod, (u[j + l] - u[j + m + l]) * y % mod)
                y = y * r0 % mod
        m >>= 1
    for i in range(l):
        u[i] = u[i] * u[i ^ l] % mod
    m = 1
    for k in range(1, x + 1):
        r0 = ir[k]
        for i in range(0, l, m << 1):
            y = 1
            for j in range(i, m + i):
                (u[j], u[j + m]) = ((u[j] + u[j + m] * y) % mod, (u[j] - u[j + m] * y) % mod)
                y = y * r0 % mod
        m <<= 1
    inv0 = pow(u[0], mod - 2, mod)
    for i in range(c):
        u[i] *= inv0
    return

def comb(n, r):
    return fact[n] * inv[r] % mod * inv[n - r] % mod if n >= r >= 0 else 0
(n, k) = map(int, input().split())
mod = 998244353
a = list(map(int, input().split()))
l = 2 * max(a) + 5
c = [0] * l
for i in a:
    c[i] += 1
b = list(map(int, input().split()))
fact = [1] * (l + 1)
for i in range(1, l + 1):
    fact[i] = i * fact[i - 1] % mod
inv = [1] * (l + 1)
inv[l] = pow(fact[l], mod - 2, mod)
for i in range(l - 1, -1, -1):
    inv[i] = (i + 1) * inv[i + 1] % mod
pow2 = [1] * (l + 1)
for i in range(1, l + 1):
    pow2[i] = 2 * pow2[i - 1] % mod
cnt = [0] * (1 << 21)
for x in b:
    (c1, c2) = (0, 0)
    for i in range(min(x, l)):
        if c[i] == 1:
            c1 += 1
        elif c[i] >= 2:
            c2 += 1
    f = [pow2[i] * comb(c1, i) % mod for i in range(c1 + 1)]
    g = [comb(2 * c2, i) for i in range(2 * c2 + 1)]
    convolve(f, g)
    m = len(f) + len(g) - 1
    j = 2 * (x + 1)
    for i in range(m):
        cnt[2 * i + j] += u[i] % mod
q = int(input())
q = list(map(int, input().split()))
ans = [cnt[i] % mod for i in q]
sys.stdout.write('\n'.join(map(str, ans)))

Solution (Python):
import sys, os, io
input = io.BytesIO(os.read(0, os.fstat(0).st_size)).readline
mod = 998244353
r = [pow(3, mod >> i, mod) for i in range(23)]
ir = [pow(i, mod - 2, mod) for i in r]
u = [0] * (1 << 23)

def convolve(lf, lg):
    c = lf + lg - 1
    x = c.bit_length()
    l = 1 << x
    for i in range(lf):
        u[i] = pow2[i] * comb(lf - 1, i) % mod
    for i in range(lf, l):
        u[i] = 0
    for i in range(lg):
        u[i ^ l] = comb(lg - 1, i)
    for i in range(lg, l):
        u[i ^ l] = 0
    m = l >> 1
    for k in range(x, 0, -1):
        r0 = r[k]
        for i in range(0, l, m << 1):
            y = 1
            for j in range(i, m + i):
                (u[j], u[j + m]) = ((u[j] + u[j + m]) % mod, (u[j] - u[j + m]) * y % mod)
                (u[j + l], u[j + m + l]) = ((u[j + l] + u[j + m + l]) % mod, (u[j + l] - u[j + m + l]) * y % mod)
                y = y * r0 % mod
        m >>= 1
    for i in range(l):
        u[i] = u[i] * u[i ^ l] % mod
    m = 1
    for k in range(1, x + 1):
        r0 = ir[k]
        for i in range(0, l, m << 1):
            y = 1
            for j in range(i, m + i):
                (u[j], u[j + m]) = ((u[j] + u[j + m] * y) % mod, (u[j] - u[j + m] * y) % mod)
                y = y * r0 % mod
        m <<= 1
    inv0 = pow(u[0], mod - 2, mod)
    for i in range(c):
        u[i] *= inv0
    return

def comb(n, r):
    return fact[n] * inv[r] % mod * inv[n - r] % mod if n >= r >= 0 else 0
(n, k) = map(int, input().split())
mod = 998244353
a = list(map(int, input().split()))
l = 2 * max(a) + 5
c = [0] * l
for i in a:
    c[i] += 1
b = list(map(int, input().split()))
fact = [1] * (l + 1)
for i in range(1, l + 1):
    fact[i] = i * fact[i - 1] % mod
inv = [1] * (l + 1)
inv[l] = pow(fact[l], mod - 2, mod)
for i in range(l - 1, -1, -1):
    inv[i] = (i + 1) * inv[i + 1] % mod
pow2 = [1] * (l + 1)
for i in range(1, l + 1):
    pow2[i] = 2 * pow2[i - 1] % mod
cnt = [0] * (1 << 21)
for x in b:
    (c1, c2) = (0, 0)
    for i in range(min(x, l)):
        if c[i] == 1:
            c1 += 1
        elif c[i] >= 2:
            c2 += 1
    convolve(c1 + 1, 2 * c2 + 1)
    j = 2 * (x + 1)
    for i in range(c1 + 2 * c2 + 1):
        cnt[2 * i + j] += u[i] % mod
q = int(input())
q = list(map(int, input().split()))
ans = [cnt[i] % mod for i in q]
sys.stdout.write('\n'.join(map(str, ans)))