競プロ初中級者100問

ABC014C AtColor

この記事で使うアルゴリズム

累積和
いもす法

はじめに

カテゴリー競プロ初中級者100問では、Qiitaにて@e869120さんがレッドコーダーが教える、競プロ・AtCoder上達のガイドライン【中級編：目指せ水色コーダー！】としてまとめられている100問をPythonで解いています。

全問題の一覧はこちらです

問題

https://atcoder.jp/contests/abc014/tasks/abc014_3

この問題では累積和を応用したいもす法というアルゴリズムで解くことができます。

方針

要素数1000000のカウンタを用意して、アンケート毎に対象の色をカウントアップすると、計算量はO(1000000n)となり間に合いません。そこでいもす法
というテクニックを利用します。

いもす法

いもす法は、配列上の連続する区間にある同じ数を足したい時に有効なテクニックです。累積和の逆のようなテクニックです。例を使って説明します。

まずは累積和のおさらいです。下記のようなリストAを考えます

1
A = [2, 3, 0, 1, 2]

Aの累積和Sは以下のようになりますね。

1
2
A = [2, 3, 0, 1, 2]
S = [2, 5, 5, 6, 8]

さて次に累積和の逆を考えます。

1
S = [0, 0, 1, 1, 1, 0, 0]

このような累積和Sを与えるAはどのようになるでしょうか。

1
2
S = [0, 0, 1, 1, 1, 0, 0]
A = [0, 0, 1, 0, 0,-1, 0]

このようになるはずです。

さて、先ほどのSのような配列が複数(S1~S3)あり、それらを足したい状況を想定しましょう。
ただし、まだ配列の形で保持しておらず、以下のように値が1となっている箇所を示す形の情報だとします。

1
2
3
S1: 2~4
S2: 3~5
S3: 2~5

この時、配列にもし展開すると下記のようになります。

1
2
3
S1 = [0, 0, 1, 1, 1, 0, 0]
S2 = [0, 0, 0, 1, 1, 1, 0]
S3 = [0, 1, 1, 1, 1, 1, 0]

これらの和Sを求めようとすると配列の数をN、各配列の要素数をKとして計算量はO(NK)となります。

一方で、先ほどのように累積和がSとなるようなAを考えます。先ほどの1となっている箇所を示す形の情報を元にするとAの作成コストは、O(N+K)です。

1
A = [0, 0, 2, 1, 0,-1,-2]

この累積和を求めると、

1
2
A = [0, 0, 2, 1, 0,-1,-2]
S = [0, 0, 2, 3, 3, 2, 0]

となります。これでより少ない計算量O(N+K)でSを求めることができました。

コード

今回の問題は、まさにいもす法の典型問題です。以下のようにシンプルに実装できます。

1
2
3
4
5
6
7
8
9
import itertools
N = int(input())
imos = [0] * (10**6 + 2)
for n in range(N):
    a, b = map(int,input().split())
    imos[a] += 1
    imos[b+1] -= 1
csum = list(itertools.accumulate(imos))
print (max(csum))

記事情報

  • 投稿日：2020年5月5日
  • 最終更新日：2020年5月8日