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:
from heapq import heappop, heappush

# use 64 bit signed integer max
NUMERIC_LIMITS = (1 << 63) - 1


class McfGraph:
  def __init__(self, n):
    self._n = n
    self.g = [[] for _ in range(n)]
    self.pos = []

  def add_edge(self, frm, to, cap, cost):
    m = len(self.pos)
    self.pos.append([frm, len(self.g[frm])])
    self.g[frm].append(self._edge(to, len(self.g[to]), cap, cost))
    self.g[to].append(self._edge(frm, len(self.g[frm]) - 1, 0, -cost))
    return m

  class edge:
    __slots__ = "frm", "to", "cap", "flow", "cost"

    def __init__(self, frm, to, cap, flow, cost):
      self.frm, self.to = frm, to
      self.cap, self.flow, self.cost = cap, flow, cost

  def get_edge(self, i):
    _e = self.g[self.pos[i][0]][self.pos[i][1]]
    _re = self.g[_e.to][_e.rev]
    return self.edge(self.pos[i][0], _e.to, _e.cap + _re.cap, _re.cap, _e.cost)

  def edges(self):
    for i in range(len(self.pos)):
      yield self.get_edge(i)

  def flow(self, s, t, flow_limit=NUMERIC_LIMITS):
    return self.slope(s, t, flow_limit)[-1]

  def dual_ref(self, s, t):
    dist = [NUMERIC_LIMITS] * self._n
    self.pv = [-1] * self._n
    self.pe = [-1] * self._n
    vis = [False] * self._n
    que = []
    dist[s] = 0
    heappush(que, (0, s))
    while que:
      _, v = heappop(que)
      if vis[v]:
        continue
      vis[v] = True
      if v == t:
        break
      for i in range(len(self.g[v])):
        e = self.g[v][i]
        if vis[e.to] or not e.cap:
          continue
        cost = e.cost - self.dual[e.to] + self.dual[v]
        if dist[e.to] - dist[v] > cost:
          dist[e.to] = dist[v] + cost
          self.pv[e.to] = v
          self.pe[e.to] = i
          heappush(que, (dist[e.to], e.to))
    if not vis[t]:
      return False
    for v in range(self._n):
      if not vis[v]:
        continue
      self.dual[v] -= dist[t] - dist[v]
    return True

  def slope(self, s, t, flow_limit=NUMERIC_LIMITS):
    self.dual = [0] * self._n

    flow = 0
    cost, prev_cost = 0, -1
    result = [(flow, cost)]
    while flow < flow_limit:
      if not self.dual_ref(s, t):
        break
      c = flow_limit - flow
      v = t
      while v != s:
        c = min(c, self.g[self.pv[v]][self.pe[v]].cap)
        v = self.pv[v]
      v = t
      while v != s:
        e = self.g[self.pv[v]][self.pe[v]]
        e.cap -= c
        self.g[v][e.rev].cap += c
        v = self.pv[v]
      d = -self.dual[s]
      flow += c
      cost += c * d
      if prev_cost == d:
        result.pop()
      result.append((flow, cost))
      prev_cost = cost
    return result

  class _edge:
    __slots__ = "to", "rev", "cap", "cost"

    def __init__(self, to, rev, cap, cost):
      self.to, self.rev, self.cap, self.cost = to, rev, cap, cost


def atcoder_practice2_e():
    import sys
    input = sys.stdin.buffer.readline

    N, K = list(map(int, input().split()))
    A = [tuple(map(int, input().split())) for _ in range(N)]
    BIG = 10 ** 9

    g = McfGraph(2 * N + 2)
    s = 2 * N
    t = 2 * N + 1

    g.add_edge(s, t, N * K, BIG)

    for i in range(N):
        g.add_edge(s, i, K, 0)
        g.add_edge(N + i, t, K, 0)
    for i in range(N):
        for j in range(N):
            g.add_edge(i, N + j, 1, BIG - A[i][j])

    result = g.flow(s, t, N * K)
    print(N * K * BIG - result[1])

    grid = [["."] * N for _ in range(N)]
    edges = g.edges()
    for e in edges:
        if e.frm == s or e.to == t or e.flow == 0:
            continue
        grid[e.frm][e.to - N] = "X"

    for g in grid:
        print("".join(g))


if __name__ == "__main__":
    atcoder_practice2_e()

Submission status: Accepted

Solution code:
# TODO: メモリリーク確認

code_mcf_graph = r"""
#define PY_SSIZE_T_CLEAN
#include <Python.h>
#include "structmember.h"

// 元のライブラリの private を剥がした

// >>> AtCoder >>>

#ifndef ATCODER_MINCOSTFLOW_HPP
#define ATCODER_MINCOSTFLOW_HPP 1

#include <algorithm>
#include <cassert>
#include <limits>
#include <queue>
#include <vector>

namespace atcoder {

template <class Cap, class Cost> struct mcf_graph {
  public:
    mcf_graph() {}
    mcf_graph(int n) : _n(n), g(n) {}

    int add_edge(int from, int to, Cap cap, Cost cost) {
        assert(0 <= from && from < _n);
        assert(0 <= to && to < _n);
        int m = int(pos.size());
        pos.push_back({from, int(g[from].size())});
        int from_id = int(g[from].size());
        int to_id = int(g[to].size());
        if (from == to) to_id++;
        g[from].push_back(_edge{to, to_id, cap, cost});
        g[to].push_back(_edge{from, from_id, 0, -cost});
        return m;
    }

    struct edge {
        int from, to;
        Cap cap, flow;
        Cost cost;
    };

    edge get_edge(int i) {
        int m = int(pos.size());
        assert(0 <= i && i < m);
        auto _e = g[pos[i].first][pos[i].second];
        auto _re = g[_e.to][_e.rev];
        return edge{
            pos[i].first, _e.to, _e.cap + _re.cap, _re.cap, _e.cost,
        };
    }
    std::vector<edge> edges() {
        int m = int(pos.size());
        std::vector<edge> result(m);
        for (int i = 0; i < m; i++) {
            result[i] = get_edge(i);
        }
        return result;
    }

    std::pair<Cap, Cost> flow(int s, int t) {
        return flow(s, t, std::numeric_limits<Cap>::max());
    }
    std::pair<Cap, Cost> flow(int s, int t, Cap flow_limit) {
        return slope(s, t, flow_limit).back();
    }
    std::vector<std::pair<Cap, Cost>> slope(int s, int t) {
        return slope(s, t, std::numeric_limits<Cap>::max());
    }
    std::vector<std::pair<Cap, Cost>> slope(int s, int t, Cap flow_limit) {
        assert(0 <= s && s < _n);
        assert(0 <= t && t < _n);
        assert(s != t);
        // variants (C = maxcost):
        // -(n-1)C <= dual[s] <= dual[i] <= dual[t] = 0
        // reduced cost (= e.cost + dual[e.from] - dual[e.to]) >= 0 for all edge
        std::vector<Cost> dual(_n, 0), dist(_n);
        std::vector<int> pv(_n), pe(_n);
        std::vector<bool> vis(_n);
        auto dual_ref = [&]() {
            std::fill(dist.begin(), dist.end(),
                      std::numeric_limits<Cost>::max());
            std::fill(pv.begin(), pv.end(), -1);
            std::fill(pe.begin(), pe.end(), -1);
            std::fill(vis.begin(), vis.end(), false);
            struct Q {
                Cost key;
                int to;
                bool operator<(Q r) const { return key > r.key; }
            };
            std::priority_queue<Q> que;
            dist[s] = 0;
            que.push(Q{0, s});
            while (!que.empty()) {
                int v = que.top().to;
                que.pop();
                if (vis[v]) continue;
                vis[v] = true;
                if (v == t) break;
                // dist[v] = shortest(s, v) + dual[s] - dual[v]
                // dist[v] >= 0 (all reduced cost are positive)
                // dist[v] <= (n-1)C
                for (int i = 0; i < int(g[v].size()); i++) {
                    auto e = g[v][i];
                    if (vis[e.to] || !e.cap) continue;
                    // |-dual[e.to] + dual[v]| <= (n-1)C
                    // cost <= C - -(n-1)C + 0 = nC
                    Cost cost = e.cost - dual[e.to] + dual[v];
                    if (dist[e.to] - dist[v] > cost) {
                        dist[e.to] = dist[v] + cost;
                        pv[e.to] = v;
                        pe[e.to] = i;
                        que.push(Q{dist[e.to], e.to});
                    }
                }
            }
            if (!vis[t]) {
                return false;
            }

            for (int v = 0; v < _n; v++) {
                if (!vis[v]) continue;
                // dual[v] = dual[v] - dist[t] + dist[v]
                //         = dual[v] - (shortest(s, t) + dual[s] - dual[t]) + (shortest(s, v) + dual[s] - dual[v])
                //         = - shortest(s, t) + dual[t] + shortest(s, v)
                //         = shortest(s, v) - shortest(s, t) >= 0 - (n-1)C
                dual[v] -= dist[t] - dist[v];
            }
            return true;
        };
        Cap flow = 0;
        Cost cost = 0, prev_cost_per_flow = -1;
        std::vector<std::pair<Cap, Cost>> result;
        result.push_back({flow, cost});
        while (flow < flow_limit) {
            if (!dual_ref()) break;
            Cap c = flow_limit - flow;
            for (int v = t; v != s; v = pv[v]) {
                c = std::min(c, g[pv[v]][pe[v]].cap);
            }
            for (int v = t; v != s; v = pv[v]) {
                auto& e = g[pv[v]][pe[v]];
                e.cap -= c;
                g[v][e.rev].cap += c;
            }
            Cost d = -dual[s];
            flow += c;
            cost += c * d;
            if (prev_cost_per_flow == d) {
                result.pop_back();
            }
            result.push_back({flow, cost});
            prev_cost_per_flow = d;
        }
        return result;
    }

//  private:
    int _n;

    struct _edge {
        int to, rev;
        Cap cap;
        Cost cost;
    };

    std::vector<std::pair<int, int>> pos;
    std::vector<std::vector<_edge>> g;
};

}  // namespace atcoder

#endif  // ATCODER_MINCOSTFLOW_HPP

// <<< AtCoder <<<

using namespace std;
using namespace atcoder;
#define PARSE_ARGS(types, ...) if(!PyArg_ParseTuple(args, types, __VA_ARGS__)) return NULL


struct MCFGraph{
    PyObject_HEAD
    mcf_graph<long long, long long>* graph;
};

struct MCFGraphEdge{
    PyObject_HEAD
    mcf_graph<long long, long long>::edge* edge;
};


extern PyTypeObject MCFGraphType;
extern PyTypeObject MCFGraphEdgeType;


// >>> MCFGraph definition >>>

static void MCFGraph_dealloc(MCFGraph* self){
    delete self->graph;
    Py_TYPE(self)->tp_free((PyObject*)self);
}
static PyObject* MCFGraph_new(PyTypeObject* type, PyObject* args, PyObject* kwds){
    return type->tp_alloc(type, 0);
}
static int MCFGraph_init(MCFGraph* self, PyObject* args){
    long n;
    if(!PyArg_ParseTuple(args, "l", &n)) return -1;
    if(n < 0 || n > (long)1e8){
        PyErr_Format(PyExc_IndexError,
            "constraint error in MCFGraph constructor (constraint: 0<=n<=1e8, got n=%d)", n);
        return -1;
    }
    self->graph = new mcf_graph<long long, long long>(n);
    return 0;
}
static PyObject* MCFGraph_add_edge(MCFGraph* self, PyObject* args){
    long from, to;
    long long cap, cost;
    PARSE_ARGS("llLL", &from, &to, &cap, &cost);
    if(from < 0 || from >= self->graph->_n || to < 0 || to >= self->graph->_n){
        PyErr_Format(PyExc_IndexError,
            "MCFGraph add_edge index out of range (n=%d, from=%d, to=%d)", self->graph->_n, from, to);
        return (PyObject*)NULL;
    }
    if(from == to){
        PyErr_Format(PyExc_IndexError, "got self-loop (from=%d, to=%d)", from, to);
        return (PyObject*)NULL;
    }
    if(cap < 0){
        PyErr_Format(PyExc_IndexError, "got negative cap (cap=%d)", cap);
        return (PyObject*)NULL;
    }
    if(cost < 0){
        PyErr_Format(PyExc_IndexError, "got negative cost (cap=%d)", cost);
        return (PyObject*)NULL;
    }
    const int res = self->graph->add_edge(from, to, cap, cost);
    return Py_BuildValue("l", res);
}
static PyObject* MCFGraph_flow(MCFGraph* self, PyObject* args){
    long s, t;
    long long flow_limit = numeric_limits<long long>::max();
    PARSE_ARGS("ll|L", &s, &t, &flow_limit);
    if(s < 0 || s >= self->graph->_n || t < 0 || t >= self->graph->_n){
        PyErr_Format(PyExc_IndexError,
            "MCFGraph flow index out of range (n=%d, s=%d, t=%d)", self->graph->_n, s, t);
        return (PyObject*)NULL;
    }
    if(s == t){
        PyErr_Format(PyExc_IndexError, "got s == t (s=%d, t=%d)", s, t);
        return (PyObject*)NULL;
    }
    const pair<long long, long long>& flow_cost = self->graph->flow(s, t, flow_limit);
    return Py_BuildValue("LL", flow_cost.first, flow_cost.second);
}
static PyObject* MCFGraph_slope(MCFGraph* self, PyObject* args){
    long s, t;
    long long flow_limit = numeric_limits<long long>::max();
    PARSE_ARGS("ll|L", &s, &t, &flow_limit);
    if(s < 0 || s >= self->graph->_n || t < 0 || t >= self->graph->_n){
        PyErr_Format(PyExc_IndexError,
            "MCFGraph slope index out of range (n=%d, s=%d, t=%d)", self->graph->_n, s, t);
        return (PyObject*)NULL;
    }
    if(s == t){
        PyErr_Format(PyExc_IndexError, "got s == t (s=%d, t=%d)", s, t);
        return (PyObject*)NULL;
    }
    const vector<pair<long long, long long>>& slope = self->graph->slope(s, t, flow_limit);
    const int siz = (int)slope.size();
    PyObject* list = PyList_New(siz);
    for(int i = 0; i < siz; i++){
        PyList_SET_ITEM(list, i, Py_BuildValue("LL", slope[i].first, slope[i].second));
    }
    return list;
}
static PyObject* MCFGraph_get_edge(MCFGraph* self, PyObject* args){
    long i;
    PARSE_ARGS("l", &i);
    const int m = (int)self->graph->pos.size();
    if(i < 0 || i >= m){
        PyErr_Format(PyExc_IndexError,
            "MCFGraph get_edge index out of range (m=%d, i=%d)", m, i);
        return (PyObject*)NULL;
    }
    MCFGraphEdge* edge = PyObject_NEW(MCFGraphEdge, &MCFGraphEdgeType);
    edge->edge = new mcf_graph<long long, long long>::edge(self->graph->get_edge(i));
    return (PyObject*)edge;
}
static PyObject* MCFGraph_edges(MCFGraph* self, PyObject* args){
    const auto& edges = self->graph->edges();
    const int m = (int)edges.size();
    PyObject* list = PyList_New(m);
    for(int i = 0; i < m; i++){
        MCFGraphEdge* edge = PyObject_NEW(MCFGraphEdge, &MCFGraphEdgeType);
        edge->edge = new mcf_graph<long long, long long>::edge(edges[i]);
        PyList_SET_ITEM(list, i, (PyObject*)edge);
    }
    return list;
}
static PyObject* MCFGraph_repr(PyObject* self){
    PyObject* edges = MCFGraph_edges((MCFGraph*)self, NULL);
    PyObject* res = PyUnicode_FromFormat("MCFGraph(%R)", edges);
    Py_DECREF(edges);
    return res;
}
static PyMethodDef MCFGraph_methods[] = {
    {"add_edge", (PyCFunction)MCFGraph_add_edge, METH_VARARGS, "Add edge"},
    {"flow", (PyCFunction)MCFGraph_flow, METH_VARARGS, "Flow"},
    {"slope", (PyCFunction)MCFGraph_slope, METH_VARARGS, "Slope"},
    {"get_edge", (PyCFunction)MCFGraph_get_edge, METH_VARARGS, "Get edge"},
    {"edges", (PyCFunction)MCFGraph_edges, METH_VARARGS, "Get edges"},
    {NULL}  /* Sentinel */
};
PyTypeObject MCFGraphType = {
    PyObject_HEAD_INIT(NULL)
    "atcoder.MCFGraph",                 /*tp_name*/
    sizeof(MCFGraph),                   /*tp_basicsize*/
    0,                                  /*tp_itemsize*/
    (destructor)MCFGraph_dealloc,       /*tp_dealloc*/
    0,                                  /*tp_print*/
    0,                                  /*tp_getattr*/
    0,                                  /*tp_setattr*/
    0,                                  /*reserved*/
    MCFGraph_repr,                      /*tp_repr*/
    0,                                  /*tp_as_number*/
    0,                                  /*tp_as_sequence*/
    0,                                  /*tp_as_mapping*/
    0,                                  /*tp_hash*/
    0,                                  /*tp_call*/
    0,                                  /*tp_str*/
    0,                                  /*tp_getattro*/
    0,                                  /*tp_setattro*/
    0,                                  /*tp_as_buffer*/
    Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE, /*tp_flags*/
    0,                                  /*tp_doc*/
    0,                                  /*tp_traverse*/
    0,                                  /*tp_clear*/
    0,                                  /*tp_richcompare*/
    0,                                  /*tp_weaklistoffset*/
    0,                                  /*tp_iter*/
    0,                                  /*tp_iternext*/
    MCFGraph_methods,                   /*tp_methods*/
    0,                                  /*tp_members*/
    0,                                  /*tp_getset*/
    0,                                  /*tp_base*/
    0,                                  /*tp_dict*/
    0,                                  /*tp_descr_get*/
    0,                                  /*tp_descr_set*/
    0,                                  /*tp_dictoffset*/
    (initproc)MCFGraph_init,            /*tp_init*/
    0,                                  /*tp_alloc*/
    MCFGraph_new,                       /*tp_new*/
    0,                                  /*tp_free*/
    0,                                  /*tp_is_gc*/
    0,                                  /*tp_bases*/
    0,                                  /*tp_mro*/
    0,                                  /*tp_cache*/
    0,                                  /*tp_subclasses*/
    0,                                  /*tp_weaklist*/
    0,                                  /*tp_del*/
    0,                                  /*tp_version_tag*/
    0,                                  /*tp_finalize*/
};

// <<< MCFGraph definition <<<


// >>> MCFGraphEdge definition >>>

static void MCFGraphEdge_dealloc(MCFGraphEdge* self){
    delete self->edge;
    Py_TYPE(self)->tp_free((PyObject*)self);
}
static PyObject* MCFGraphEdge_new(PyTypeObject* type, PyObject* args, PyObject* kwds){
    return type->tp_alloc(type, 0);
}
static int MCFGraphEdge_init(MCFGraphEdge* self, PyObject* args){
    int from, to;
    long long cap, flow, cost;
    if(!PyArg_ParseTuple(args, "llLLL", &from, &to, &cap, &flow, &cost)) return -1;
    self->edge = new mcf_graph<long long, long long>::edge(mcf_graph<long long, long long>::edge{from, to, cap, flow, cost});
    return 0;
}
static PyObject* MCFGraphEdge_get_from(MCFGraphEdge* self, PyObject* args){
    return PyLong_FromLong(self->edge->from);
}
static PyObject* MCFGraphEdge_get_to(MCFGraphEdge* self, PyObject* args){
    return PyLong_FromLong(self->edge->to);
}
static PyObject* MCFGraphEdge_get_flow(MCFGraphEdge* self, PyObject* args){
    return PyLong_FromLongLong(self->edge->flow);
}
static PyObject* MCFGraphEdge_get_cap(MCFGraphEdge* self, PyObject* args){
    return PyLong_FromLongLong(self->edge->cap);
}
static PyObject* MCFGraphEdge_get_cost(MCFGraphEdge* self, PyObject* args){
    return PyLong_FromLongLong(self->edge->cost);
}
static PyObject* MCFGraphEdge_repr(PyObject* self){
    MCFGraphEdge* self_ = (MCFGraphEdge*)self;
    PyObject* res = PyUnicode_FromFormat("MCFGraphEdge(%2d -> %2d, %2lld / %2lld, cost = %2lld)",
        self_->edge->from, self_->edge->to, self_->edge->flow, self_->edge->cap, self_->edge->cost);
    return res;
}
PyGetSetDef MCFGraphEdge_getsets[] = {
    {"from_", (getter)MCFGraphEdge_get_from, NULL, NULL, NULL},
    {"to", (getter)MCFGraphEdge_get_to, NULL, NULL, NULL},
    {"flow", (getter)MCFGraphEdge_get_flow, NULL, NULL, NULL},
    {"cap", (getter)MCFGraphEdge_get_cap, NULL, NULL, NULL},
    {"cost", (getter)MCFGraphEdge_get_cost, NULL, NULL, NULL},
    {NULL}
};
PyTypeObject MCFGraphEdgeType = {
    PyObject_HEAD_INIT(NULL)
    "atcoder.MCFGraphEdge",             /*tp_name*/
    sizeof(MCFGraphEdge),               /*tp_basicsize*/
    0,                                  /*tp_itemsize*/
    (destructor)MCFGraphEdge_dealloc,   /*tp_dealloc*/
    0,                                  /*tp_print*/
    0,                                  /*tp_getattr*/
    0,                                  /*tp_setattr*/
    0,                                  /*reserved*/
    MCFGraphEdge_repr,                  /*tp_repr*/
    0,                                  /*tp_as_number*/
    0,                                  /*tp_as_sequence*/
    0,                                  /*tp_as_mapping*/
    0,                                  /*tp_hash*/
    0,                                  /*tp_call*/
    0,                                  /*tp_str*/
    0,                                  /*tp_getattro*/
    0,                                  /*tp_setattro*/
    0,                                  /*tp_as_buffer*/
    Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE, /*tp_flags*/
    0,                                  /*tp_doc*/
    0,                                  /*tp_traverse*/
    0,                                  /*tp_clear*/
    0,                                  /*tp_richcompare*/
    0,                                  /*tp_weaklistoffset*/
    0,                                  /*tp_iter*/
    0,                                  /*tp_iternext*/
    0,                                  /*tp_methods*/
    0,                                  /*tp_members*/
    MCFGraphEdge_getsets,               /*tp_getset*/
    0,                                  /*tp_base*/
    0,                                  /*tp_dict*/
    0,                                  /*tp_descr_get*/
    0,                                  /*tp_descr_set*/
    0,                                  /*tp_dictoffset*/
    (initproc)MCFGraphEdge_init,        /*tp_init*/
    0,                                  /*tp_alloc*/
    MCFGraphEdge_new,                   /*tp_new*/
    0,                                  /*tp_free*/
    0,                                  /*tp_is_gc*/
    0,                                  /*tp_bases*/
    0,                                  /*tp_mro*/
    0,                                  /*tp_cache*/
    0,                                  /*tp_subclasses*/
    0,                                  /*tp_weaklist*/
    0,                                  /*tp_del*/
    0,                                  /*tp_version_tag*/
    0,                                  /*tp_finalize*/
};

// <<< MCFGraphEdge definition <<<


static PyModuleDef atcodermodule = {
    PyModuleDef_HEAD_INIT,
    "atcoder",
    NULL,
    -1,
};

PyMODINIT_FUNC PyInit_atcoder(void)
{
    PyObject* m;
    if(PyType_Ready(&MCFGraphType) < 0) return NULL;
    if(PyType_Ready(&MCFGraphEdgeType) < 0) return NULL;

    m = PyModule_Create(&atcodermodule);
    if(m == NULL) return NULL;

    Py_INCREF(&MCFGraphType);
    if (PyModule_AddObject(m, "MCFGraph", (PyObject*)&MCFGraphType) < 0) {
        Py_DECREF(&MCFGraphType);
        Py_DECREF(m);
        return NULL;
    }
    
    Py_INCREF(&MCFGraphEdgeType);
    if (PyModule_AddObject(m, "MCFGraphEdge", (PyObject*)&MCFGraphEdgeType) < 0) {
        Py_DECREF(&MCFGraphEdgeType);
        Py_DECREF(m);
        return NULL;
    }
    return m;
}
"""
code_mcf_graph_setup = r"""
from distutils.core import setup, Extension
module = Extension(
    "atcoder",
    sources=["mcf_graph.cpp"],
    extra_compile_args=["-O3", "-march=native", "-std=c++14"]
)
setup(
    name="atcoder-library",
    version="0.0.1",
    description="wrapper for atcoder library",
    ext_modules=[module]
)
"""

import os
import sys

if sys.argv[-1] == "ONLINE_JUDGE" or os.getcwd() != "/imojudge/sandbox":
    with open("mcf_graph.cpp", "w") as f:
        f.write(code_mcf_graph)
    with open("mcf_graph_setup.py", "w") as f:
        f.write(code_mcf_graph_setup)
    os.system(f"{sys.executable} mcf_graph_setup.py build_ext --inplace")

from atcoder import MCFGraph, MCFGraphEdge


def main():
    N, K = map(int, input().split())
    g = MCFGraph(2 * N + 2)
    BIG = 10 ** 9
    s = 2 * N
    t = 2 * N + 1
    g.add_edge(s, t, N * K, BIG)
    for i in range(N):
        g.add_edge(s, i, K, 0)
        g.add_edge(N + i, t, K, 0)
        A = list(map(int, input().split()))
        for j, a in enumerate(A):
            g.add_edge(i, N + j, 1, BIG - a)
    res = g.flow(s, t, N * K)[1]
    print(N * K * BIG - res)
    edges = g.edges()
    grid = [["."] * N for _ in range(N)]
    for e in edges:
        if e.from_ == s or e.to == t or e.flow == 0:
            continue
        grid[e.from_][e.to - N] = "X"
    for gr in grid:
        print("".join(gr))

main()

Submission status: Accepted

Solution code:
import sys

sys.setrecursionlimit(10**7)
def MI(): return map(int,sys.stdin.readline().rstrip().split())
def LI(): return list(map(int,sys.stdin.readline().rstrip().split()))  #空白あり


class MinCostFlow:
    def __init__(self, n):
        self.n = n
        self.G = [[] for _ in range(n)]

    def addEdge(self, f, t, cap, cost):
        # [to, cap, cost, rev]
        self.G[f].append([t, cap, cost, len(self.G[t])])
        self.G[t].append([f, 0, -cost, len(self.G[f])-1])

    def minCostFlow(self, s, t, f):
        n = self.n
        G = self.G
        prevv = [0]*n; preve = [0]*n
        INF = 10**18

        res = 0
        while f:
            dist = [INF]*n
            dist[s] = 0
            update = 1
            while update:
                update = 0
                for v in range(n):
                    if dist[v] == INF:
                        continue
                    gv = G[v]
                    for i in range(len(gv)):
                        to, cap, cost, rev = gv[i]
                        if cap > 0 and dist[v] + cost < dist[to]:
                            dist[to] = dist[v] + cost
                            prevv[to] = v; preve[to] = i
                            update = 1
            if dist[t] == INF:
                return -1

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


N,K = MI()
A = [LI() for _ in range(N)]
max = 10**13

MCF = MinCostFlow(2*N+2)

s,t = 2*N,2*N+1

MCF.addEdge(s,t,N*K,max)

for i in range(N):
    MCF.addEdge(s,i,K,0)
    MCF.addEdge(i+N,t,K,0)
    for j in range(N):
        MCF.addEdge(i,j+N,1,max-A[i][j])

print(N*K*max-MCF.minCostFlow(s,t,N*K))

ANS = [['.']*N for _ in range(N)]

for i in range(N):
    for to,cap,cost,_ in MCF.G[i]:
        if cap == 0 and N <= to < 2*N:
            ANS[i][to-N] = 'X'

for i in range(N):
    print(''.join(ANS[i]))

Submission status: Accepted

Solution code:
import sys
import numpy as np
import networkx as nx

read = sys.stdin.buffer.read
readline = sys.stdin.buffer.readline
readlines = sys.stdin.buffer.readlines

X = 1 << 30

N, K = map(int, readline().split())
A = np.array(read().split(), np.int64).reshape(N, N)

G = nx.DiGraph()
G.add_nodes_from(range(N + N + 2))
source, sink = N + N, N + N + 1
G.nodes[source]['demand'] = -N * K
G.nodes[sink]['demand'] = N * K
for i in range(N):
    G.add_edge(source, i, capacity=K, weight=0)
    G.add_edge(N + i, sink, capacity=K, weight=0)
    G.add_edge(source, sink, capacity=N*K, weight=X)  # K個ずつ使わなくてもよい
for i in range(N):
    for j in range(N):
        G.add_edge(i, N + j, capacity=1, weight=X - A[i, j])

mcf = nx.min_cost_flow(G)
mcf_cost = nx.cost_of_flow(
    G,
    mcf,
)

mat = np.zeros((N, N), np.int8)
for i in range(N):
    for to, f in mcf[i].items():
        mat[i, to - N] = f

print(N * K * X - mcf_cost)
for row in mat:
    print(''.join(map('.X'.__getitem__, row)))