日本レジストリサービス（JPRS）プログラミングコンテスト2025#1（AtCoder Beginner Contest 392）F,G問題メモ

日本レジストリサービス（JPRS）プログラミングコンテスト2025#1（AtCoder Beginner Contest 392）

比較的全完しやすい回だった。

F - Insert

• 空の配列 $A$ があります。 $i=1,2,\ldots,N$ の順に以下の操作を行います。
• 数 $i$ を、$A$ の前から $P_i$ 番目の位置になるように挿入する。
  • より正確には、「$A$ の $P_i-1$ 項目まで」「 $i$ 」「$A$ の $P_i$ 項目以降」をこの順に連結したもので $A$ を置き換える。

• 全ての操作を終えた後の $A$ を出力してください。

• $1 \leq N \leq 5\times 10^5$
• $1 \leq P_i \leq i$
• 入力は全て整数である

以下の例で考えてみる。

i   1 2 3 4 5 6
P   1 2 2 1 3 2

■A の作られ過程
1   1
2   1 2
3   1 3 2
4   4 1 3 2
5   4 1 5 3 2
6   4 6 1 5 3 2

わかりやすいところでは、 例えば最後の $i=6$ は、挿入後、（最後なので当然）自身の前に割り込む要素がないので、 最終状態でも $P_6=2$ の示す通り、先頭から2番目の位置に置かれている。

$i=5$ は、挿入後、$i=6$ という1つの要素に割り込まれるので、最終状態では $P_5+1=4$ 番目に置かれている。
基本的には $i$ が大きい方が、自身の処理後に割り込む候補となる数が少ないので、考えやすそうである。

$i$ の大きい方から処理することを考えると、以下のような問題に言い換えることができる。

• 長さ $N$ の配列 $A$ があり、各要素は null で初期化されています。
• $i=N,\ldots,2,1$ の順に以下の操作を行います。
  • $A$ の中の null のうち、先頭から $P_i$ 番目に $i$ を置く。

i   1 2 3 4 5 6
P   1 2 2 1 3 2

6   _ 6 _ _ _ _
5   _ 6 _ 5 _ _
4   4 6 _ 5 _ _
3   4 6 _ 5 3 _
2   4 6 _ 5 3 2
1   4 6 1 5 3 2

SortedSetなど、「$i$ 番目の要素を取得する」「要素 $x$ を消す」という2つの処理を行えるデータ構造で実現できる。
最初、$S=\{1,2,...,N\}$ としておいて、$i$ ごとに以下をすればよい。

• $S$ の残っている要素の $P_i$ 番目を取得する。$k$ とする。
• $A_k=i$ とする。
• $S$ から $k$ を削除する。

Python3

from sortedcontainers import SortedSet


def solve(n, ppp):
    ss = SortedSet(range(n))
    ans = [0] * n
    for i in range(n - 1, -1, -1):
        p = ppp[i] - 1
        k = ss[p]
        ans[k] = i + 1
        ss.remove(k)
    return ans


n = int(input())
ppp = list(map(int, input().split()))
ans = solve(n, ppp)
print(*ans)

G - Fine Triplets

• 整数 $A,B,C$ ( $A \lt B \lt C$ ) が $B-A = C-B$ を満たす時、 $(A,B,C)$ を 良い三つ組 と呼びます。
• $N$ 要素からなる正整数の集合 $S=\{ S_1,S_2,\dots,S_N \}$ が与えられるので、 $A,B,C \in S$ なる良い三つ組の個数を求めてください。

• 入力は全て整数
• $1 \le N \le 10^6$
• $1 \le S_i \le 10^6$
• $S$ の要素は相異なる

仮に以下が分かっていたとすると、

• $P[i]:=$「$A+C=i$ となる $(A,C)$ ペアの個数」

$B$ を固定したときの良い三つ組の個数は、$P[2B]$ で取得できる。

そして $P$ は畳み込みで求められる。

• $F[i]:=i$ が $S$ に含まれるなら1、それ以外は0

$F$ 同士を畳み込んだ数列を $G$ とすると、$G[i]=2P[i]+\alpha$ となる。
$\alpha$ は、$i=2B$ となる $B$ が $S$ に含まれる場合 $1$、それ以外は $0$ となる変数とする。

（$(A,C),(C,A)$ で2回数えられる分と、$(B,B)$ 自身が1回数えられている分）

ここから $P$ が求められる。

Python3

import numpy as np


def convolve(f, g):
    fft_len = 1
    true_len = len(f) + len(g) - 1
    while fft_len < true_len:
        fft_len <<= 1

    Ff = np.fft.rfft(f, fft_len)
    Fg = np.fft.rfft(g, fft_len)
    Fh = Ff * Fg

    h = np.fft.irfft(Fh, fft_len)
    h = np.rint(h).astype(np.int64)

    return h[:true_len]


def solve(n, sss):
    f = np.bincount(sss)
    g = convolve(f, f).tolist()
    ans = 0
    for s in sss:
        ans += g[s * 2] - 1
    ans //= 2
    return ans


n = int(input())
sss = list(map(int, input().split()))
ans = solve(n, sss)
print(ans)