【AtCoder】ABC 429 D - On AtCoder Conference

2025-10-31

AtCoder/ABC/D問題 AtCoder/ABC/425点問題 AtCoder/緑Diff AtCoder/NoviSteps/1Q AtCoder/累積和 AtCoder/二分探索/lower_bound AtCoder/円環は2周分考える AtCoder/データ構造/連想配列(map)競技プログラミング

https://atcoder.jp/contests/abc429/tasks/abc429_d

atcoder.jp

実行時間制限: 2 sec / メモリ制限: 1024 MiB / Difficulty: 903 / NoviSteps: 1Q / 配点: 425 点

問題概要

$1$ 周の長さが $M$ である池があり、その周上に $1$ つの小屋と $N$ 人の人が立っている。

実数 $x$ $(0\leq x < M)$ について、小屋から時計回りに $x$ だけ進んだところを地点 $x$ と定義する。 $i$ 番目の人は、地点 $A_i$ にいて、複数の人が同じ場所に立っていることもある。

また、 $N$ 以下の整数 $C$ が与えられ、 $i=0,1,\ldots,M-1$ について、 $X_i$ を次のように定義する。

高橋君は地点 $(i+0.5)$ からスタートして、時計回りに動き始める。
高橋君は出会った人数の合計が $C$ 未満である限り時計回りに動き続け、 $C$ 以上になったら止まる。
高橋君が止まるまでに出会った人数を $X_i$ とする。ここで、高橋君が止まった地点に複数の人がいる場合、そこにいた人はすべて出会った人として数えられ、特に $X_i$ は $C$ より大きくなる可能性がある。

$i=0,1,\ldots,M-1$ にわたる $X_i$ の総和 $\sum_{i=0}^{M-1} X_i$ を求めよ。

制約

$1\leq N\leq 5\times 10^5$
$1\leq M\leq 10^{12}$
$0\leq A_i\leq M-1$
$1\leq C\leq N$
入力はすべて整数

考察

同じ値を取る $X_i$ のグループ分け

まず、各 $X_i$ の値は、高橋君のスタート地点 $i+0.5$ が人のいる地点を通過するときにしか変化しないことに注目しよう。これを利用することで、考えるべき $X_i$ の値をスタート地点の位置で以下のようにグループ分けすることができる。

なお、以下では人のいるユニークな地点を昇順にソートしたものを $p_0, p_1, \ldots, p_{K-1}$ とする。

$p_{j} \leq i < p_{j+1} - 1\: (0 \leq j < K-1)$ のパターン

このとき、高橋君がスタートしてから最初に出会うのは地点 $p_{j+1}$ にいる人たちである。

したがって、この範囲のスタート地点 $i$ に対して、 $X_i$ の値はすべて同じになる。

この範囲に存在する $i$ の個数は、 $(p_{j+1} - 1) - p_{j} + 1 = p_{j+1} - p_{j}$ 個である。

$p_{K-1} \leq i \leq M - 1$ (スタート地点が最後の人のいる地点より右側) のパターン

このとき、高橋君がスタートしてから最初に出会うのは、小屋を越えた(あるいは小屋と同じ)地点 $p_{0}$ にいる人たちである。

したがって、この範囲のスタート地点 $i$ に対して、 $X_i$ の値はすべて同じになる。

この範囲に存在する $i$ の個数は、 $(M - 1) - p_{K-1} + 1 = M - p_{K-1}$ 個である。

$0 \leq i \leq p_0 - 1$ (スタート地点が最初の人のいる地点より左側) のパターン

このとき、高橋君がスタートしてから最初に出会うのは、地点 $p_{0}$ にいる人たちである。

したがって、この範囲のスタート地点 $i$ に対して、 $X_i$ の値はすべて同じになる。

この範囲に存在する $i$ の個数は、 $p_{0}$ 個である。

ここまでの考察をまとめると、求めるべき $\sum_{i=0}^{M-1} X_i$ の値は、以下のように $K$ 個のグループに分けて計算することができる。

\sum_{i=0}^{M-1} X_i = \sum_{j=0}^{K-2} (p_{j+1} - p_{j}) X_{p_{j}} + (M - p_{K-1}) X_{p_{K-1}} + p_{0} X_{p_{0}}

ただし、ここでの $X_{p_j}$ は、区間 $[p_j, p_{j+1}-1]$ に属するすべての $X_i$ で共通の値を表す。

各グループの $X_{p_j}$ の計算

次に、各グループにおける $X_{p_j}$ の値を計算する方法を考えよう。スタート地点から $C$ 人に出会うまでを愚直にシミュレーションすると計算量が $O(N)$ となり、全体では $O(NK)$ となってしまって間に合わない。

そこで、円環が出てくる問題の典型であるデータを 2 倍に増やして線形化する手法を用いる。

具体的には、今考えている $p_j$ の位置とその地点に立っている人の数のペアを $(p_j, cj)$ と表すことにし、これらを $2$ 倍に増やして以下のように並べる。

(p_0, c_0), (p_1, c_1), \ldots, (p_{K-1}, c_{K-1}), (p_0 + M, c_0), (p_1 + M, c_1), \ldots, (p_{K-1} + M, c_{K-1})

その上で、この配列の 1 番目の要素に対して累積和テーブルを作成し、「スタート地点から数えて出会った人数が $C$ 人以上になる最小のインデックス」を二分探索で求めることで、各 $X_{p_j}$ の値を $O(\log K)$ で計算することができる。

以上の処理を実装することで、全体の計算量を $O(N \log N)$ として解くことができ、これは本問の制約下で十分に高速である。

実装例

実装では、 $p_j$ のリストを pos、その 2 倍に増やしたリストを pos2 とし、累積和テーブルを ps としている。

1.#include <bits/stdc++.h>
2.using namespace std;
3. 
4.#define rep(i, start, end) for (auto i = (start); (i) < (end); (i)++)
5. 
6.using ll = long long;
7. 
8.// ======================================== //
9. 
10.int main()
11.{
12.    ll N, M, C;
13.    cin >> N >> M >> C;
14.    map<ll, int> mp;
15.    rep(i, 0, N)
16.    {
17.        ll A;
18.        cin >> A;
19.        mp[A]++;
20.    }
21. 
22.    vector<pair<ll, int>> pos;
23.    for (auto &[key, val] : mp)
24.    {
25.        pos.push_back({key, val});
26.    }
27. 
28.    int K = pos.size();
29.    vector<pair<ll, int>> pos2 = pos;
30.    rep(i, 0, K) pos2.push_back({pos[i].first + M, pos[i].second});
31. 
32.    vector<ll> ps(2 * K + 1, 0);
33.    rep(i, 0, 2 * K) ps[i + 1] = ps[i] + pos2[i].second;
34. 
35.    auto get_x = [&](int start) -> ll
36.    {
37.        ll target_num = ps[start] + C;
38. 
39.        auto it = lower_bound(ps.begin() + start + 1, ps.end(), target_num);
40.        int target_idx = distance(ps.begin(), it) - 1;
41. 
42.        return ps[target_idx + 1] - ps[start];
43.    };
44. 
45.    ll ans = 0;
46.    rep(j, 0, K - 1)
47.    {
48.        ll q = pos[j + 1].first - pos[j].first;
49.        ans += q * get_x(j + 1);
50.    }
51. 
52.    ll q_end = M - pos[K - 1].first;
53.    ans += q_end * get_x(K);
54. 
55.    ll q_start = pos[0].first;
56.    ans += q_start * get_x(0);
57. 
58.    cout << ans << endl;
59. 
60.    return 0;
61.}