ABC360 D問題（Ghost Ants）を解く

AtCoder が提供しているABC（AtCoder Beginner Contest）360 D問題をC++とPythonで解いてみました。ABC360は、2024年6月30日21:00に実施されました。

AtCoder の紹介はこちらに、プログラミングの方針はこちらに記事があります。

D問題 Ghost Ants（Difficulty : 159）

問題の詳細は、リンク先をご覧ください。

ABC360 D問題 Ghost Ants

ある範囲に含まれている数字の個数を二分探索で求めます。AtCoder Problems による Difficulty は 159 でした。

解答案

蟻を負の向きに進む toL と、正の向きに進む toR に分けます。toR の要素に対して、toR[i] ≦ toL[j] ≦ toR[i] + 2 × T となる toL の個数を求める問題になります。

toL をソートして、二分探索を使うことにより計算が間に合います。

C++ プログラム例（ABC360D）

二分探索の補足をします。

• lower_bound は、指定した値以上の最小値のイテレータを返します。指定した値未満の要素の個数が分かります。
• upper_bound は、指定した値を超える最小値のイテレータを返します。指定した値以下の要素の個数が分かります。
• 値が見つからない場合は、end() が指すイテレータを返しますが、引き算をして個数に変換することができます（31ー33行目）。
  • はじめに全ての蟻は相異なる座標にいるため、31行目は、upper_bound を使っても結果は同じとなります。

以下が、C++プログラムです。

#include <bits/stdc++.h>
using namespace std;

typedef long long int ll;
int main()
{
	int n;
	cin >> n;
	ll t;
	cin >> t;
	string s;
	cin >> s;
	vector<ll> x(n);
	for (int i = 0; i < n; ++i) {
		cin >> x[i];
	}

	vector<ll> toL;
	vector<ll> toR;
	for (int i = 0; i < n; ++i) {
		if (s[i] == '0') {
			toL.push_back(x[i]);
		} else {
			toR.push_back(x[i]);
		}
	}
	sort(toL.begin(), toL.end());

	ll result = 0;
	for (int i = 0; i < toR.size(); ++i) {
		auto it1 = lower_bound(toL.begin(), toL.end(), toR[i]);
		auto it2 = upper_bound(toL.begin(), toL.end(), toR[i] + 2LL * t);
		result += it2 - it1;
	}

	cout << result << endl;

	return 0;
}

AC（Accepted＝正しいプログラム）と判定されました。

Python プログラム例（ABC360D）

Python版も基本的な考え方は同じです。二分探索は、bisect を使いました。bisect_left が lower_bound に、bisect_right が upper_bound に相当します。以下がプログラムです。

"""AtCoder Beginner Contest 360 D"""
import bisect

n, t = map(int, input().split())
s = input()
x = list(map(int, input().split()))

toL = []
toR = []
for i in range(n):
    if s[i] == '0':
        toL.append(x[i])
    else:
        toR.append(x[i])
toL.sort()

result = 0
for e in toR:
    result += bisect.bisect_right(toL, e + 2 * t) \
            - bisect.bisect_left(toL, e)

print(result)

こちらも「AC」と判定されました。

最後に

「蟻本」（「プログラミングコンテストチャレンジブック」第2版 秋葉拓哉、岩田陽一、北川宣稔著、マイナビ 2012年）で近い問題が解説されています。コンテストの問題では、すれ違いで向きは変わりません（Ghostなのでしょう）。ただし、蟻本にあるように蟻の向きが反対方向になっても、すれ違う回数は同じです。結果としての軌跡が一致するからです。

コンテストでは、upper_bound で書くべきところを lower_bound で書いてしまい、WAが取れずに時間切れでした。残念です。

引き続き ABC の問題を紹介していきます。