Problem:


Score : 100 points

Problem Statement

You are given a grid of N rows and M columns. The square at the i-th row and j-th column will be denoted as (i,j).
A nonnegative integer A_{i,j} is written for each square (i,j).

You choose some of the squares so that each row and column contains at most K chosen squares.
Under this constraint, calculate the maximum value of the sum of the integers written on the chosen squares.
Additionally, calculate a way to choose squares that acheives the maximum.

Constraints

1 \leq N \leq 50

1 \leq K \leq N

0 \leq A_{i,j} \leq 10^9

All values in Input are integer.

Input

Input is given from Standard Input in the following format:

N K
A_{1,1} A_{1,2} \cdots A_{1,N}
A_{2,1} A_{2,2} \cdots A_{2,N}
\vdots
A_{N,1} A_{N,2} \cdots A_{N,N}

Output

On the first line, print the maximum value of the sum of the integers written on the chosen squares.

On the next N lines, print a way that achieves the maximum.

Precisely, output the strings t_1,t_2,\cdots,t_N, that satisfies t_{i,j}=X if you choose (i,j) and t_{i,j}=. otherwise.

You may print any way to choose squares that maximizes the sum.

Sample Input 1

3 1
5 3 2
1 4 8
7 6 9

Sample Output 1

19
X..
..X
.X.

Sample Input 2

3 2
10 10 1
10 10 1
1 1 10

Sample Output 2

50
XX.
XX.
..X



Sample input:
3 1
5 3 2
1 4 8
7 6 9


Sample output:
19
X..
..X
.X.

Submission status: Accepted

Solution code:
import os
import sys
from heapq import heappop, heappush

import numpy as np


def solve(n, k, aaa):
    MINCOSTFLOW_LINKS = []
    MINCOSTFLOW_LINK_INDICES = []
    MINCOSTFLOW_CURRENT_LINK_INDEX = []
    INF = 10 ** 10

    def mincostflow_init(m):
        """ m:辺数(逆辺除く) """
        MINCOSTFLOW_LINKS.append(np.zeros((m * 2, 5), np.int64))
        MINCOSTFLOW_CURRENT_LINK_INDEX.append(0)
        return len(MINCOSTFLOW_LINKS) - 1

    def mincostflow_add_link(ins, frm, to, capacity, cost):
        links = MINCOSTFLOW_LINKS[ins]
        li = MINCOSTFLOW_CURRENT_LINK_INDEX[ins]
        links[li] = (frm, to, capacity, cost, li + 1)
        links[li + 1] = (to, frm, 0, -cost, li)
        MINCOSTFLOW_CURRENT_LINK_INDEX[ins] += 2

    def mincostflow_fix_link(ins, n):
        links = MINCOSTFLOW_LINKS[ins]
        sort_idx = np.argsort(links[:, 0])
        links[links[sort_idx, 4], 4] = np.arange(sort_idx.size)
        links = links[sort_idx]

        MINCOSTFLOW_LINK_INDICES.append(np.searchsorted(links[:, 0], np.arange(n + 1)))
        MINCOSTFLOW_LINKS[ins] = links

    def mincostflow_flow(ins, s, t, quantity):
        links = MINCOSTFLOW_LINKS[ins]
        link_indices = MINCOSTFLOW_LINK_INDICES[ins]
        n = len(links)
        res = 0
        potentials = np.zeros(n, dtype=np.int64)
        dist = np.full(n, INF, dtype=np.int64)
        prev_v = np.full(n, -1, dtype=np.int64)
        prev_e = np.full(n, -1, dtype=np.int64)

        while quantity:
            dist.fill(INF)
            dist[s] = 0
            que = [(0, s)]

            while que:
                total_cost, v = heappop(que)
                if dist[v] < total_cost:
                    continue
                for li in range(link_indices[v], link_indices[v + 1]):
                    _, u, cap, cost, _ = links[li]
                    new_cost = dist[v] + potentials[v] - potentials[u] + cost
                    if cap > 0 and new_cost < dist[u]:
                        dist[u] = new_cost
                        prev_v[u] = v
                        prev_e[u] = li
                        heappush(que, (new_cost, u))

            # Cannot flow required quantity
            if dist[t] == INF:
                return -1

            potentials += dist

            cur_flow = quantity
            v = t
            while v != s:
                cur_flow = min(cur_flow, links[prev_e[v]][2])
                v = prev_v[v]
            quantity -= cur_flow
            res += cur_flow * potentials[t]

            v = t
            while v != s:
                link = links[prev_e[v]]
                link[2] -= cur_flow
                links[link[4]][2] += cur_flow
                v = prev_v[v]

        return res

    ins = mincostflow_init(n ** 2 + n * 2 + 1)
    s = n + n
    t = s + 1

    mincostflow_add_link(ins, s, t, n * k, INF)
    for i in range(n):
        mincostflow_add_link(ins, s, i, k, 0)
        mincostflow_add_link(ins, i + n, t, k, 0)
        for j in range(n):
            mincostflow_add_link(ins, i, j + n, 1, INF - aaa[i][j])
    mincostflow_fix_link(ins, n * 2 + 2)

    ans = INF * (n * k) - mincostflow_flow(ins, s, t, n * k)

    matrix = np.zeros((n, n), dtype=np.int8)
    links = MINCOSTFLOW_LINKS[ins]
    link_indices = MINCOSTFLOW_LINK_INDICES[ins]
    for i in range(n):
        for li in range(link_indices[i], link_indices[i + 1]):
            link = links[li]
            if link[1] == t or link[2] == 1 or link[3] <= 0:
                continue
            j = link[1] - n
            matrix[i][j] = 1

    return ans, matrix


SIGNATURE = '(i8,i8,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)

n, k = map(int, sys.stdin.readline().split())
aaa = np.fromstring(sys.stdin.read(), dtype=np.int64, sep=' ').reshape(n, n)
ans, matrix = solve(n, k, aaa)

sys.stdout.buffer.write((str(ans) + '\n').encode())
np.savetxt(sys.stdout.buffer, np.vectorize('.X'.__getitem__)(matrix), fmt='%s', delimiter='')

Submission status: Accepted

Solution code:
import networkx as nx
from networkx.algorithms.flow import network_simplex

MAX = 10**9
n, k = map(int, input().split())
a = [list(map(int, input().split())) for _ in range(n)]
s, t = 2 * n, 2 * n + 1

G = nx.DiGraph()
G.add_node(s, demand=-n*k)
G.add_node(t, demand=n*k)
G.add_edge(s, t, weight=MAX, capacity=n*k)

for i in range(n):
	G.add_edge(s, i, weight=0, capacity=k)
	G.add_edge(n + i, t, weight=0, capacity=k)
	for j in range(n):
		G.add_edge(i, n + j, weight=MAX-a[i][j], capacity=1)

flow_cost, flow_dict = network_simplex(G)
print(n * k * MAX - flow_cost)

ans = [["." for _ in range(n)] for _ in range(n)]
for fr in flow_dict.keys():
	if 0 <= fr < n:
		for to, flow in flow_dict[fr].items():
			if flow == 1:
				ans[fr][to - n] = "X"

print(*["".join(x) for x in ans], sep="\n")

Submission status: Accepted

Solution code:
import sys

sys.path.append("/home/contestant/.local/lib/python3.8/site-packages")

import networkx as nx

N, K = [int(x) for x in input().split()]
A = [[int(x) for x in input().split()] for i in range(N)]


graph = nx.DiGraph()

SOURCE = 2 * N
SINK = SOURCE + 1

graph.add_node(SOURCE, demand=-K * N)
for i in range(N):
    graph.add_node(i)  # row
    graph.add_node(i + N)  # col
graph.add_node(SINK, demand=K * N)


for r in range(N):
    graph.add_edge(SOURCE, r, capacity=K, weight=0)
for r in range(N):
    for c in range(N):
        graph.add_edge(r, c + N, capacity=1, weight=-A[r][c])  # neg to maximize
for r in range(N):
    graph.add_edge(r + N, SINK, capacity=K, weight=0)

graph.add_edge(SOURCE, SINK, capacity=N * K, weight=0)  # slack


flowDict = nx.algorithms.flow.min_cost_flow(graph)
total = 0
grid = [["." for c in range(N)] for r in range(N)]
for r in range(N):
    for c in range(N):
        if flowDict[r][c + N]:
            grid[r][c] = "X"
            total += A[r][c]

print(total)
for row in grid:
    print("".join(row))

Submission status: Accepted

Solution code:
import networkx as nx

N, K = [int(x) for x in input().split()]
A = [[int(x) for x in input().split()] for i in range(N)]


graph = nx.DiGraph()

SOURCE = 2 * N
SINK = SOURCE + 1

graph.add_node(SOURCE, demand=-K * N)
for i in range(N):
    graph.add_node(i)  # row
    graph.add_node(i + N)  # col
graph.add_node(SINK, demand=K * N)


for r in range(N):
    graph.add_edge(SOURCE, r, capacity=K, weight=0)
for r in range(N):
    for c in range(N):
        graph.add_edge(r, c + N, capacity=1, weight=-A[r][c])  # neg to maximize
for r in range(N):
    graph.add_edge(r + N, SINK, capacity=K, weight=0)

graph.add_edge(SOURCE, SINK, capacity=N * K, weight=0)  # slack


flowDict = nx.algorithms.flow.min_cost_flow(graph)
total = 0
grid = [["." for c in range(N)] for r in range(N)]
for r in range(N):
    for c in range(N):
        if flowDict[r][c + N]:
            grid[r][c] = "X"
            total += A[r][c]

print(total)
for row in grid:
    print("".join(row))

Submission status: Accepted

Solution code:
import sys
readline = sys.stdin.readline
from heapq import heappop as hpp, heappush as hp
class MinCostFlowwithDijkstra:
    INF = 1<<60
    
    def __init__(self, N):
        self.N = N
        self.Edge = [[] for _ in range(N)]
    
    def add_edge(self, st, en, cap, cost):
        self.Edge[st].append([en, cap, cost, len(self.Edge[en])])
        self.Edge[en].append([st, 0, -cost, len(self.Edge[st])-1])
    
    def get_mf(self, so, si, fl):
        N = self.N
        INF = self.INF
        res = 0
        Pot = [0]*N
        geta = N
        
        
        prv = [None]*N
        prenum = [None]*N
        while fl:
            dist = [INF]*N
            dist[so] = 0
            Q = [so]
            
            while Q:
                cost, vn = divmod(hpp(Q), geta)
                if dist[vn] < cost:
                    continue
                
                for enum in range(len(self.Edge[vn])):
                    vf, cap, cost, _ = self.Edge[vn][enum]
                    cc = dist[vn] + cost - Pot[vn] + Pot[vf]
                    if cap > 0 and dist[vf] > cc:
                        dist[vf] = cc
                        prv[vf] = vn
                        prenum[vf] = enum
                        hp(Q, cc*geta + vf)
            
            if dist[si] == INF:
                return -1
            
            for i in range(N):
                Pot[i] -= dist[i]
            
            cfl = fl
            vf = si
            while vf != so:
                cfl = min(cfl, self.Edge[prv[vf]][prenum[vf]][1])
                vf = prv[vf]
            
            fl -= cfl
            res -= cfl*Pot[si]
            vf = si
            while vf != so:
                e = self.Edge[prv[vf]][prenum[vf]]
                e[1] -= cfl
                self.Edge[vf][e[3]][1] += cfl
                vf = prv[vf]
        return res


N, K = map(int, readline().split())

MAX = 10**9+7

A = [list(map(int, readline().split())) for _ in range(N)]

T = MinCostFlowwithDijkstra(2*N+2)
st = 2*N
en = st+1

T.add_edge(st, en, MAX, MAX)
for i in range(N):
    T.add_edge(st, i, K, 0)
    T.add_edge(N+i, en, K, 0)

for i in range(N):
    for j in range(N):
        T.add_edge(i, N+j, 1, MAX-A[i][j])

ff = N//K * K*K + (N%K)**2 

print(MAX*MAX - T.get_mf(st, en, MAX))
dot = ord('.')
cross = ord('X')
Ans = [[dot]*N for _ in range(N)]

cnt = 0
for j in range(N):
    for vf, cap, _, _ in T.Edge[N+j]:
        if cap == 1 and vf != en:
            Ans[vf][j] = cross
print('\n'.join(''.join(map(chr, a)) for a in Ans))

Submission status: Accepted

Solution code:
import sys
input = lambda : sys.stdin.readline().rstrip()
sys.setrecursionlimit(max(1000, 10**9))
write = lambda x: sys.stdout.write(x+"\n")


n,k = list(map(int, input().split()))
a = [list(map(int, input().split())) for _ in range(n)]
import networkx as nx
ans = (-float("inf"), None)
f = n*k
g = nx.DiGraph()
for i in range(2*n):
    g.add_node(i)
g.add_node(2*n, demand=-f)
g.add_node(2*n+1, demand=f)
for i in range(n):
    g.add_edge(2*n, i, capacity=k, weight=0)
    g.add_edge(i+n, 2*n+1, capacity=k, weight=0)
g.add_edge(2*n, 2*n+1, capacity=f, weight=0)
for i in range(n):
    for j in range(n):
        g.add_edge(i, j+n, capacity=1, weight=-a[i][j])
val, d = nx.network_simplex(g)
vals = [["."]*n for _ in range(n)]
for i in range(n):
    for j in range(n):
        if d[i][j+n]==1:
            vals[i][j] = "X"
if -val>ans[0]:
    ans = (-val, vals)
print(ans[0])
print("\n".join("".join(ans[1][i]) for i in range(n)))

Submission status: Accepted

Solution code:
import os
import sys
from heapq import heappop, heappush

import numpy as np


def solve(n, k, aaa):
    MINCOSTFLOW_LINKS = []
    INF = 10 ** 10

    def mincostflow_init(n):
        lst = [[0]]
        lst.clear()
        MINCOSTFLOW_LINKS.append([lst.copy() for _ in range(n)])
        return len(MINCOSTFLOW_LINKS) - 1

    def mincostflow_add_link(ins, frm, to, capacity, cost):
        links = MINCOSTFLOW_LINKS[ins]
        links[frm].append([to, capacity, cost, len(links[to])])
        links[to].append([frm, 0, -cost, len(links[frm]) - 1])

    def mincostflow_flow(ins, s, t, f):
        links = MINCOSTFLOW_LINKS[ins]
        n = len(links)
        res = 0
        hierarchy = np.zeros(n, dtype=np.int64)
        dist = np.full(n, INF, dtype=np.int64)
        prev_v = np.full(n, -1, dtype=np.int64)
        prev_e = np.full(n, -1, dtype=np.int64)

        while f:
            dist.fill(INF)
            dist[s] = 0
            que = [(0, s)]

            while que:
                total_cost, v = heappop(que)
                # print(total_cost, v)
                if dist[v] < total_cost:
                    continue
                dv = dist[v] + hierarchy[v]
                for i, (u, cap, cost, _) in enumerate(links[v]):
                    new_cost = dv + cost - hierarchy[u]
                    if cap > 0 and new_cost < dist[u]:
                        dist[u] = new_cost
                        prev_v[u] = v
                        prev_e[u] = i
                        heappush(que, (new_cost, u))

            # Cannot flow quantity f
            if dist[t] == INF:
                return -1

            hierarchy += dist

            cur_flow = f
            v = t
            while v != s:
                cur_flow = min(cur_flow, links[prev_v[v]][prev_e[v]][1])
                v = prev_v[v]
            f -= cur_flow
            res += cur_flow * hierarchy[t]

            v = t
            while v != s:
                link = links[prev_v[v]][prev_e[v]]
                link[1] -= cur_flow
                links[v][link[3]][1] += cur_flow
                v = prev_v[v]

        return res

    ins = mincostflow_init(n * 2 + 2)
    s = n + n
    t = s + 1

    mincostflow_add_link(ins, s, t, n * k, INF)
    for i in range(n):
        mincostflow_add_link(ins, s, i, k, 0)
        mincostflow_add_link(ins, i + n, t, k, 0)
        for j in range(n):
            mincostflow_add_link(ins, i, j + n, 1, INF - aaa[i][j])
    ans = INF * (n * k) - mincostflow_flow(ins, s, t, n * k)

    matrix = np.zeros((n, n), dtype=np.int8)
    links = MINCOSTFLOW_LINKS[ins]
    for i in range(n):
        for link in links[i]:
            if link[0] == t or link[1] == 1 or link[2] <= 0:
                continue
            j = link[0] - n
            matrix[i][j] = 1

    return ans, matrix


SIGNATURE = '(i8,i8,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)

n, k = map(int, sys.stdin.readline().split())
aaa = np.fromstring(sys.stdin.read(), dtype=np.int64, sep=' ').reshape(n, n)
ans, matrix = solve(n, k, aaa)

print(ans)
CHARS = '.X'.__getitem__
for row in matrix:
    print(''.join(map(CHARS, row)))

Submission status: Accepted

Solution code:
import heapq


class PrimalDual:
    def __init__(self, n: int):
        """頂点数をnとする。"""
        self.INF = 10 ** 9 + 7
        self.n = n
        self.graph = [[] for _ in range(n)]

    def add_edge(self, _from: int, to: int, capacity: int, cost: int):
        """辺を追加
        1. _fromからtoへ向かう容量capacity、単位コストcostの辺をグラフに追加する。
        2. toから_fromへ向かう容量0、単位コスト-costの辺をグラフに追加する。
        """
        forward = [to, capacity, cost, None]
        forward[3] = backward = [_from, 0, -cost, forward]
        self.graph[_from].append(forward)
        self.graph[to].append(backward)

    def min_cost_flow(self, s: int, t: int, f: int) -> int:
        """s-tパス上に流量fを流すときの最小費用流を求める。
        計算量: O(|F||E|log|V|)
        """
        res = 0
        potential = [0] * self.n
        prv_v = [0] * self.n
        prv_e = [None] * self.n
        while f > 0:
            # ポテンシャルを用いたダイクストラ法
            dist = [self.INF] * self.n
            dist[s] = 0
            q = [(0, s)]  # q = [(sからのコスト, 現在地)]
            while q:
                cost, _from = heapq.heappop(q)
                if dist[_from] < cost:
                    continue
                for edge in self.graph[_from]:
                    to, capacity, cost, _ = edge
                    p_diff = potential[_from] - potential[to]
                    if capacity > 0 and dist[_from] + cost + p_diff < dist[to]:
                        dist[to] = dist[_from] + cost + p_diff
                        prv_v[to] = _from
                        prv_e[to] = edge
                        heapq.heappush(q, (dist[to], to))

            if dist[t] == self.INF:
                return -1
            for i in range(self.n):
                if dist[i] != self.INF:
                    potential[i] += dist[i]
            d = f
            v = t
            while v != s:
                d = min(d, prv_e[v][1])
                v = prv_v[v]
            f -= d
            res += potential[t] * d
            v = t
            while v != s:
                edge = prv_e[v]
                edge[1] -= d
                edge[3][1] += d
                v = prv_v[v]
        return res


n, k = map(int, input().split())
a = [list(map(int, input().split())) for i in range(n)]
INF = 10 ** 18


pd = PrimalDual(n + n + 2)

s = n + n
t = n + n + 1

pd.add_edge(s, t, INF, 0)
for i in range(n):
    pd.add_edge(s, n + i, k, 0)
    pd.add_edge(i, t, k, 0)
for i in range(n):
    for j in range(n):
        pd.add_edge(n + i, j, 1, -a[i][j])

ans = -pd.min_cost_flow(s, t, INF)

res = [["." for i in range(n)] for i in range(n)]
for i in range(n):
    for j, cap, cost, _ in pd.graph[n + i]:
        if 0 <= j <= n and cap == 0:
            res[i][j] = "X"

print(ans)
for r in res:
    print("".join(r))

Submission status: Accepted

Solution code:
from networkx import*
N, K = map(int, input().split())
A = [[int(a) for a in input().split()] for _ in range(N)]
G = MultiDiGraph()
s = 2 * N * (N + 1)
t = s + 1
G.add_node(s, demand = -K*N)
G.add_node(t, demand = K*N)
G.add_nodes_from([i for i in range(s - 1)], demand = 0)
G.add_weighted_edges_from([(s, t, 10 ** 9)], capacity = N * K)
G.add_weighted_edges_from([(s, i, 0) for i in range(2*N*N, N*(2*N+1))], capacity = K)
G.add_weighted_edges_from([(i, t, 0) for i in range(N*(2*N+1), 2*N*(N+1))], capacity = K)
for i in range(N):
    G.add_weighted_edges_from([(2*N*N+i, i*N+j, 0) for j in range(N)], capacity = 1)
    G.add_weighted_edges_from([(N*N+i*N+j, N*(2*N+1)+j, 0) for j in range(N)], capacity = 1)
for i in range(N):
    G.add_weighted_edges_from([(i*N+j, N*N+i*N+j, 10**9-A[i][j]) for j in range(N)], capacity = 1)
    
mcf = min_cost_flow(G)
ans = 0
B = [["."] * N for _ in range(N)]
for a in mcf:
    for b in mcf[a]:
        v = mcf[a][b][0]
        if v and b == a + N*N and 0 <= a < N*N:
            i, j = a // N, a % N
            ans += v * A[i][j]
            B[i][j] = "X"
print(ans)
for b in B:
    print("".join(b))

Submission status: Runtime Error

Solution code:
# https://atcoder.jp/contests/practice2/tasks/practice2_e

import sys

from atcoder.mincostflow import MCFGraph


def main() -> None:
    n, k = map(int, sys.stdin.readline().split())
    s = n * 2
    t = s + 1
    g = MCFGraph(t + 1)
    grid = [list(map(int, sys.stdin.readline().split())) for _ in range(n)]
    big = 10 ** 9
    edges = list([] for _ in range(n))
    for i in range(n):
        g.add_edge(s, i, k, 0)
        g.add_edge(i + n, t, k, 0)
        for j in range(n):
            edges[i].append(g.add_edge(i, j + n, 1, big - grid[i][j]))
    g.add_edge(s, t, n * k, big)

    cost = g.flow(s, t, n * k)[1]
    print(big * n * k - cost)

    result = list(['.'] * n for _ in range(n))
    for i in range(n):
        for j in range(n):
            if g.get_edge(edges[i][j]).flow > 0:
                result[i][j] = 'X'

    for s in result:
        print("".join(s))


if __name__ == '__main__':
    main()

Submission status: Runtime Error

Solution code:
import sys

def input(): return sys.stdin.readline().strip()
def list2d(a, b, c): return [[c] * b for i in range(a)]
def list3d(a, b, c, d): return [[[d] * c for j in range(b)] for i in range(a)]
def list4d(a, b, c, d, e): return [[[[e] * d for j in range(c)] for j in range(b)] for i in range(a)]
def ceil(x, y=1): return int(-(-x // y))
def INT(): return int(input())
def MAP(): return map(int, input().split())
def LIST(N=None): return list(MAP()) if N is None else [INT() for i in range(N)]
def Yes(): print('Yes')
def No(): print('No')
def YES(): print('YES')
def NO(): print('NO')
sys.setrecursionlimit(10**9)
INF = 10**19
MOD = 10**9 + 7
EPS = 10**-10

class MinCostFlow:
    """ 最小費用流(ベルマンフォード版、負コストに対応可)：O(F*E*V) """

    INF = 10**18

    def __init__(self, N):
        self.N = N
        self.G = [[] for i in range(N)]
        self.pos = []
    
    def copy(self):
        res = MinCostFlow(self.N)
        res.G = [[a[:] for a in b] for b in self.G]
        res.pos = [a[:] for a in self.pos]
        return res

    def add_edge(self, fr, to, cap, cost):
        self.pos.append((fr, len(self.G[fr])))
        self.G[fr].append([to, cap, cost, len(self.G[to])])
        self.G[to].append([fr, 0, -cost, len(self.G[fr])-1])

    def flow(self, s, t, f):

        N = self.N; G = self.G
        INF = MinCostFlow.INF

        res = 0
        prv_v = [0] * N
        prv_e = [0] * N

        while f:
            dist = [INF] * N
            dist[s] = 0
            update = True

            while update:
                update = False
                for v in range(N):
                    if dist[v] == INF:
                        continue
                    for i, (to, cap, cost, _) in enumerate(G[v]):
                        if cap > 0 and dist[to] > dist[v] + cost:
                            dist[to] = dist[v] + cost
                            prv_v[to] = v; prv_e[to] = i
                            update = True
            if dist[t] == INF:
                return INF

            d = f; v = t
            while v != s:
                d = min(d, G[prv_v[v]][prv_e[v]][1])
                v = prv_v[v]
            f -= d
            res += d * dist[t]
            v = t
            while v != s:
                e = G[prv_v[v]][prv_e[v]]
                e[1] -= d
                G[v][e[3]][1] += d
                v = prv_v[v]
        return res

    def get_edge(self, i):
        e = self.G[self.pos[i][0]][self.pos[i][1]]
        re = self.G[e[0]][e[3]]
        # from, to, cap, flow, cost
        return (self.pos[i][0], e[0], e[1]+re[1], re[1], e[2])

    def get_edges(self):
        M = len(self.pos)
        res = []
        for i in range(M):
            res.append(self.get_edge(i))
        return res

def trisearch_min(lo, hi, func):
    """ 三分探索 """

    while lo+2 < hi:
        m1 = (lo*2+hi) // 3
        m2 = (lo+hi*2) // 3
        res1 = func(m1)
        res2 = func(m2)
        if res1 <= res2:
            hi = m2
        else:
            lo = m1
    return m1, m2

N, K = MAP()
grid = [[] for i in range(N)]
for i in range(N):
    grid[i] = LIST()

mcf = MinCostFlow(N*2+2)
s = N*2
t = N*2 + 1
for i in range(N):
    mcf.add_edge(s, i, K, 0)
    mcf.add_edge(N+i, t, K, 0)
    for j in range(N):
        mcf.add_edge(i, N+j, 1, -grid[i][j])

def check(x):
    global mcf

    tmp = mcf.copy()
    res = mcf.flow(s, t, x)
    mcf = tmp
    return res

m1, m2 = trisearch_min(0, N*N+1, check)
x = m1 if check(m1) <= check(m2) else m2

mx = -mcf.flow(s, t, x)
ans = [['.' for j in range(N)] for i in range(N)]
edges = mcf.get_edges()
for fr, to, cap, f, c in edges:
    if fr != s and to != t and f:
        h = fr
        w = to - N
        ans[h][w] = 'X'
print(mx)
for i in range(N):
    print(''.join(ans[i]))

Submission status: Runtime Error

Solution code:
import numpy as np

n,m=map(int,input().split())
mod=998244353
f = np.array([0]*(2**19))
g = np.array([0]*(2**19))
f[0:n]=list(map(int,input().split()))
g[0:m]=list(map(int,input().split()))
fft_len = 2**20  # |f|+|g|-1 以上の2の冪数
 
# f,gをフーリエ変換
Ff = np.fft.rfft(f, fft_len)
#print(Ff)
Fg = np.fft.rfft(g, fft_len)
#print(Fg)
 
# 各要素について掛け合わせる
Ffg = Ff * Fg
#print(Ffg)
 
# 逆変換後、整数に丸める
fg = np.fft.irfft(Ffg, fft_len)
fg = np.rint(fg)
for i in range(n+m-1):
  print(int(fg[i])%mod,end=" ")