AtCoder が提供しているABC（AtCoder Beginner Contest）371 E問題をC++とPythonで解いてみました。ABC371は、2024年9月14日21:00に実施されました。
AtCoder の紹介はこちらに、プログラミングの方針はこちらに記事があります。
E問題 I Hate Sigma Problems（Difficulty : 981）
問題の詳細は、リンク先をご覧ください。
ABC371 E問題 I Hate Sigma Problems
特定の数字が含まれている区間の数の合計を求めます。AtCoder Problems による Difficulty は 981 でした。
解答案
前提の知識
入力例1：数列 {1, 2, 2} の全ての区間はいくつあるでしょうか。
{1}、{1, 2}、{1, 2, 2}、{2}、{2, 2}、{2} の6個あります。数字を ‘○’、両端と数字の間を ‘|’ と表現すると、”|○|○|○|” となります。’|’ が4個あります。すべての区間は、この4個の ‘|’ から2つ選択することと等しくなります。従って、すべての区間の数は、$_4C_2 = (4 \times 3) / (2 \times 1) = 6$ となります。
一般的に $N$ 個の数字からなる数列のすべての区間の数は、$_{n+1}C_2 = (n + 1) n / 2$ となります。
問題を読み解く
問題は、「すべての区間に含まれる数字の種類の数の合計を求めよ」となります。これは、以下のように書き換えます。
すべての区間に含まれる数字の種類の数の合計
= $1$を含む区間の数 $+$ $2$を含む区間の数 $+ \cdots +$ $N$ を含む区間の数
となります。
入力例2：数列 {5, 4, 2, 2, 3, 2, 4, 4, 1} で 4 を含む区間の数を考えてみます。方針は、4を含まない区間の数を求めて、全体から引くことにします。
すべての区間の数は、$_{9+1}C_2 = 45$ 通りあります。ここで4を ‘○’、4以外を ‘×’、と表現するとこの数列は、
× ○ × × × × ○ ○ ×
となります。ここで「4を含まない区間の数」を求めます。’○’ を経由しないため、最初の1個の ‘×’ から成る区間の数は、$_{1+1}C_2 = 1$ 通りとなります。次に4個の ‘×’ から成る区間の数は、$_{4+1}C_2 = 10$ 通りとなります。最後の1個の ‘×’ から成る区間の数は、$_{1+1}C_2 = 1$ 通りです。
4を含む区間の数 = $45 -(1 + 10 + 1) = 33$
となります。すべての数について、同じように求めて合計すれば求める解となります。
C++ プログラム例（ABC371E）
上記の方針で実装します。
- 数字の場所を配列 b に格納します（17行目）。
- それぞれの数字を含む区間の数を求めます。
- 最初の添え字 $-1$ と、最後の添え字 $n$ として計算します。
- 全体の区間の数を加えます（24行目）。
- その数字ではない長さから、その数字を含まない区間の数を引きます（25ー28行目）。
以下が、C++プログラムです。
#include <bits/stdc++.h>
using namespace std;
typedef long long int ll;
ll nc2(ll n) { return n * (n - 1LL) / 2LL; }
int main()
{
int n;
cin >> n;
vector<int> a(n);
vector<vector<int>> b(n);
for (int i = 0; i < n; ++i) {
cin >> a[i];
--a[i];
b[a[i]].push_back(i);
}
ll result = 0;
for (int i = 0; i < n; ++i) {
b[i].push_back(n);
int pre = -1;
result += nc2(n + 1);
for (int e : b[i]) {
result -= nc2(e - pre);
pre = e;
}
}
cout << result << endl;
return 0;
}AC（Accepted＝正しいプログラム）と判定されました。
Python プログラム例（ABC371E）
Python版も基本的な考え方は同じです。以下がプログラムです。
"""AtCoder Beginner Contest 371 E"""
def nc2(n):
return n * (n - 1) // 2
n = int(input())
a = list(map(int, input().split()))
b = [[] for _ in range(n)]
for i in range(n):
a[i] -= 1
b[a[i]].append(i)
result = 0
for i in range(n):
b[i].append(n)
pre = -1
result += nc2(n + 1)
for e in b[i]:
result -= nc2(e - pre)
pre = e
print(result)こちらも「AC」と判定されました。
最後に
この問題の解説は、公式解説動画を参考にしました。勉強になりました。
引き続き ABC の問題を紹介していきます。
