パナソニックプログラミングコンテスト2021(AtCoder Beginner Contest 231) F,H問題メモ

$i=j$ の場合、$(A_i,B_i)$ と $(A_j,B_j)$ が異なる場合、同じ場合、 の3つを別々に数える冗長なコードとなっている。

最初に $(A_i,B_i)$ 毎に個数を数え、同じものはまとめて処理すれば、よりスッキリ書ける。

from collections import Counter
from operator import add


class FenwickTreeInjectable:
    def __init__(self, n, identity_factory, func):
        self.size = n
        self.tree = [identity_factory() for _ in range(n + 1)]
        self.func = func
        self.idf = identity_factory
        self.depth = n.bit_length()

    def add(self, i, x):
        i += 1
        tree = self.tree
        func = self.func
        while i <= self.size:
            tree[i] = func(tree[i], x)
            i += i & -i

    def sum(self, i):
        i += 1
        s = self.idf()
        tree = self.tree
        func = self.func
        while i > 0:
            s = func(s, tree[i])
            i -= i & -i
        return s

    def lower_bound(self, x, lt):
        """
        累積和がx以上になる最小のindexと、その直前までの累積和

        :param lt: lt(a, b) で a < b ならTrueを返す関数
        """
        total = self.idf()
        pos = 0
        tree = self.tree
        func = self.func
        for i in range(self.depth, -1, -1):
            k = pos + (1 << i)
            if k > self.size:
                continue
            new_total = func(total, tree[k])
            if lt(new_total, x):
                total = new_total
                pos += 1 << i
        return pos + 1, total

    def debug_print(self):
        prev = 0
        arr = []
        for i in range(self.size):
            curr = self.sum(i)
            arr.append(curr - prev)
            prev = curr
        print(arr)


n = int(input())
aaa = list(map(int, input().split()))
bbb = list(map(int, input().split()))
points = []
compress = set()
for a, b in zip(aaa, bbb):
    points.append((a, -b))
    compress.add(b)
points.sort()

b_dict = {b: i + 1 for i, b in enumerate(sorted(compress))}

fwt = FenwickTreeInjectable(n, int, add)
ans = n
for i, (a, b) in enumerate(points):
    bi = b_dict[-b]
    lo = fwt.sum(bi - 1)
    ans += i - lo
    fwt.add(bi, 1)

cnt = Counter(points)
for _, c in cnt.most_common():
    if c == 1:
        break
    ans += (c * c - c) // 2

print(ans)

上記の説明と $P,Q ↔ S,T$ の意味が逆なので注意

class MinimumCostFlow:
    def __init__(self, n):
        self.n = n
        self.edges = [[] for _ in range(n)]
        self.INF = 10 ** 18

    def add_edge(self, from_, to, capacity, cost):
        es, et = self.edges[from_], self.edges[to]
        ls, lt = len(es), len(et)
        es.append([capacity, to, cost, lt])
        et.append([0, from_, -cost, ls])

    def _trace_back(self, sink, predecessors):
        p = predecessors[sink]
        while p is not None:
            v, i = p
            yield self.edges[v][i]
            p = predecessors[v]

    def flow(self, source, sink, required_flow):
        n = self.n
        edges = self.edges
        INF = self.INF
        trace_back = self._trace_back

        res = 0
        while required_flow:
            dist = [INF] * n
            dist[source] = 0
            predecessors = [None] * n

            while True:
                updated = False
                for v in range(n):
                    if dist[v] == INF:
                        continue
                    for i, (remain, target, cost, _) in enumerate(edges[v]):
                        new_dist = dist[v] + cost
                        if remain and dist[target] > new_dist:
                            dist[target] = new_dist
                            predecessors[target] = (v, i)
                            updated = True
                if not updated:
                    break

            if dist[sink] == INF:
                return -1

            aug = min(required_flow, min(e[0] for e in trace_back(sink, predecessors)))
            required_flow -= aug
            res += aug * dist[sink]
            for e in trace_back(sink, predecessors):
                remain, target, cost, idx = e
                e[0] -= aug
                edges[target][idx][0] += aug
        return res


h, w, n = map(int, input().split())
mcf = MinimumCostFlow(h + w + 4)
p = h + w
q = p + 1
s = q + 1
t = s + 1
INF = 1 << 60
for _ in range(n):
    a, b, c = map(int, input().split())
    a -= 1
    b -= 1
    b += h
    mcf.add_edge(a, b, 1, c)
for i in range(h):
    mcf.add_edge(s, i, 1, 0)
    mcf.add_edge(p, i, INF, 0)
for j in range(h, h + w):
    mcf.add_edge(j, t, 1, 0)
    mcf.add_edge(j, q, INF, 0)
mcf.add_edge(p, q, INF, 0)
mcf.add_edge(s, p, w, 0)
mcf.add_edge(q, t, h, 0)
ans = mcf.flow(s, t, h + w)
print(ans)