ABC416-D を Python で解いた

ABC416 D - Match, Mod, Minimize 2
https://atcoder.jp/contests/abc416/tasks/abc416_d

なんとか今回もD問題まで解けました。しかしかなり遠回りなことをしていたようです。解説を読んで感心しました。それはそれとして、自分が解いた方法を書きます。

コンテスト中の動き

C問題までを16分で片付け、D問題を44分かけて解くことができました。E問題は計算量を落とす方法がわからず断念しました……（解説によれば計算量は大丈夫だったみたいですが）。

考察

1つめの例を見ながら全てのパターンを書き出しました。長さ N の数列 A を並び替えるので N! 通りのパターンがありますね。試しにそれぞれ足しているうちに「それぞれの和の mod M の総和を求めるが、求まった総和は M で割らない」というポイントに気づきました。もし総和の mod M を最小にする問題だったら大変でしたね。僕には解けないです。

そこからは「足してあまりが 0 になる = 和が M の倍数になる A[i] と B[i] の組を探して消していく」みたいなことも考えました。そのうち、A[i] と B[i] の和がとる範囲が最大でも $ 2M - 1 $ であることに気づきました。そこで、

• 和が 0 か M になる組み合わせ
• 和が 1 か M+1 になる組み合わせ
• 和が 2 か M+2 になる組み合わせ
• ......
• 和が M-1 か 2M-1 になる組み合わせ

を順番に見つけて消していけばいいのでは？と考えました。ただし、1つずつ探していたら計算が間に合いません。

（経緯は忘れましたが）そのうちに思いついたのが、それぞれ A[i] に対して最適な B[i] を見つけていく方針です。ある A[i] について A[i] + B[i] の取り得る値は A[i] + 0 から A[i] + M - 1 までです。この中で M で割ったあまりを小さくするには、B[i] の中でも最小のものを A[i] に足すか、A[i] + B[i] が M 以上かつなるべく小さくするようにするかのどちらかが良さそうです。というわけで両者を比較して小さい方を採用します。そして「M 以上かつなるべく小さくする」ものを探すは明らかに二分探索が有効ですね。B をあらかじめソートしておいてから二分探索です。

実装

• ペアが揃ったら順次消していく
• B はソートされていてほしい

というわけで B の中身は SortedDict に格納します。key は M で割ったあまり、value はその個数です。二分探索には SortedDict にある bisect_left 関数を用いる……と言いたいところですが key だけを取り出して使いたいので SD.keys() として key を取り出し、それに bisect.bisect_left 関数を使います。

from sortedcontainers import SortedDict
import bisect
T = int(input())
ans = []
for _ in range(T):
  N, M = map(int, input().split())
  A = list(map(int, input().split()))
  B = list(map(int, input().split()))
  SD = SortedDict()
  for b in B:
    bb = b % M
    if bb not in SD:
      SD[bb] = 0
    SD[bb] += 1
  tmp_ans = 0
  A.sort()
  
  for i in range(N):
    # 今残っている B の中で最小の値を取得し、A に足した結果を num1 とする。
    key1 = SD.peekitem(0)[0] 
    num1 = (A[i] + key1) % M
    # 今残っている B の中で A と足して M 以上になるもので最小のものを取得し、A に足した結果を num2 とする。
    sd_keys = SD.keys() 
    idx = bisect.bisect_left(sd_keys, M - A[i])
    if idx == len(SD):  # 該当するものがなければ自動的に最小の B を採用する。
      SD[key1] -= 1
      if SD[key1] == 0:
        del SD[key1]
      tmp_ans += num1
    else: # num1, num2 を比較して小さい方を採用する。
      key2 = SD.peekitem(idx)[0]
      num2 = (A[i] + key2) % M
      if num1 <= num2:
        SD[key1] -= 1
        if SD[key1] == 0:
          del SD[key1]
        tmp_ans += num1
      else:
        SD[key2] -= 1
        if SD[key2] == 0:
          del SD[key2]
        tmp_ans += num2
  ans.append(tmp_ans)

for an in ans:
  print(an)

なお、実際には B に含まれる数値は全て M 未満なのでわざわざ M で割ったあまりにする必要はなかったみたいです。また、Aをあらかじめソートしておかないとなぜか TLE が 2 つ出てしまいました（コンテスト中の提出時にはつけていましたが不要と判断して外して試しました）。2027ms という微妙な値なのでほんの少しの差なのでしょうが……まあ感覚的には A も小さい方から見た方が早く済みそうではありますね。なんとか正解していますが微妙に怪しい解法だったかもしれません。

実装2

追記です。上記の理由により SortedDict を使うのをやめて配列Bを SortedList に格納するようにしました。コードが大幅に短縮し、速度も上がりました。Aのソートがなくても余裕で回ります。さっきの実装はAのソートありで1556ms, こちらはソートなしで 1152ms です。

from sortedcontainers import SortedList
import bisect
T = int(input())
ans = []
for _ in range(T):
  N, M = map(int, input().split())
  A = list(map(int, input().split()))
  B = SortedList(map(int, input().split()))
  tmp_ans = 0
  for i in range(N):
    # 今残っている B の中で最小の値を取得し、A に足した結果を num1 とする。
    num1 = (A[i] + B[0]) % M
    # 今残っている B の中で A と足して M 以上になるもので最小のものを取得し、A に足した結果を num2 とする。
    idx = bisect.bisect_left(B, M - A[i])
    if idx == len(B):  # 該当するものがなければ自動的に最小の B を採用する。
      B.pop(0)
      tmp_ans += num1
    else: # num1, num2 を比較して小さい方を採用する。
      num2 = (A[i] + B[idx]) % M
      if num1 <= num2:
        B.pop(0)
        tmp_ans += num1
      else:
        B.pop(idx)
        tmp_ans += num2
  ans.append(tmp_ans)

for an in ans:
  print(an)

公式の解法

解説を読んで感心しましたが、よくよく考えると僕も去年この考え方を使って解いた問題がありました。ABC353 C - Sigma Problem ですね。

• A の範囲が 10^8 未満だから A[i] + A[j] を 10^8 で割った商は 0 か 1 のどちらかしかない
• 全体の合計は数式ですぐに求められる
• 全体の合計から 10^8 を超える組の個数分を引けばあまりの総和が求められる
• 一方を固定し、ソートされた他方の数列から和が 10^8 以上になる位置を二分探索で求める

今回の問題は 10^8 を M に変えただけです。

総和を 10^8 で割ったあまりを求めるわけではないことに注意してください。

という点も今回の問題と同じです。もちろん問題の核になる考え方が共通しているだけで、内容は今回の方がかなりややこしくなってはいます。