JOI 2023/2024 2次予選 問題2 解答解説

この解説は、予選Aランク(本戦出場)を目指す人むけの解説になります。
(対象外:プログラムを習い始めたばかりの人、すでにAランクに到達している人)

そのため、説明をあえて省略した部分や、プログラムに冗長な部分があります。あらかじめご了承ください。
(万が一、記載内容に誤りや不備がありましたら、ご一報いただけますと幸いです)

B – 買い物 2 (Shopping 2)

問題 https://www2.ioi-jp.org/joi/2023/2024-yo2/2024-yo2-t2.html
AtCoder https://atcoder.jp/contests/joi2024yo2/tasks/joi2024_yo2_b

まずは、問題文を理解する(15点解答)

まずは問題文を正しく理解するところから始めましょう。
「N個の商品が左から順番に並んでいて、L番目からR番目までの商品を1個ずつ買ったときの合計金額を求めよ」という問題を、お客さんの数Q人分行ってほしい。ただし、お客さんが来た日Tと商品の種類番号Aが一致する場合は、その商品は半額になる。という設定です。
これを素直にプログラムすると以下のようになります。

n, m, q = map(int, input().split())
p = [0] * n
a = [0] * n
for i in range(n):
    p[i], a[i] = map(int, input().split())

for i in range(q):
    t, l, r = map(int, input().split())
    l -= 1
    r -= 1
    ans = 0
    for j in range(l, r + 1):
        if a[j] == t:
            ans += p[j] // 2
        else:
            ans += p[j]
    print(ans)

このプログラムの計算量はO(NQ)です。小課題1のみ合格となり、15点が得られます。

M日分の累積和を作っておく(47点解答)

「L番目からR番目までの値の合計の計算をQ回分繰り返す」という問題のパターンでは、累積和というテクニックが非常に有効です。

《参考》累積和を何も考えずに書けるようにする!
https://qiita.com/drken/items/56a6b68edef8fc605821

今回の問題では、T日目には商品番号AがTと同じ商品が半額という条件があるので、通常の累積和に加えてひと工夫が必要になります。
例えば、セールはM日間続き、1<=Tk<=Mなので、M日分の累積和を事前に作ることが考えられます。 以下はこの方針で作成したプログラムです。

n, m, q = map(int, input().split())
pa = [[0 for _ in range(n + 1)] for _ in range(m + 1)]
for i in range(n):
    p, a = map(int, input().split())
    for j in range(1, m + 1):
        if j == a:
            pa[j][i + 1] = pa[j][i] + p // 2
        else:
            pa[j][i + 1] = pa[j][i] + p

for _ in range(q):
    t, l, r = map(int, input().split())
    print(pa[t][r] - pa[t][l - 1])

このプログラムの計算量はO(NM+Q)です。小課題1,2,3に合格し、47点が得られます。

二分探索を用いた満点解法

満点まであと少しです。
M日分の累積和を作るところがボトルネックになっているので、その部分を改良しましょう。
先程のプログラムでは、1つの商品に対してセール日以外の累積和の計算も必要になるため無駄が多いプログラムになっています。

そこで、

  • 種類数分、同じ種類番号をもつものだけで累積和をとっておく(累積和はm_sumとします)
  • 累積和とは別に、同じ種類番号をもつ商品の商品番号をリストに管理しておく(m_indexesとします)
  • T日目に、種類番号がTである商品のうち商品番号がLからRの範囲にはいるものをm_indexesから二分探索でしらべ、その値をもとにm_sumから購入した半額対象商品の合計金額を調べる。
  • 全商品の累積和(p_sum)をもとに、購入した商品の定価から、半額対象商品の合計の半分を引いたものが、そのお客さんの購入金額となる

この方針で作成したプログラムが以下になります。

# 二分探索
def lower_bound(temp_list, key):
    left = -1
    right = len(temp_list)
    while right - left > 1:
        mid = (left + right) // 2
        if temp_list[mid] >= key:
            right = mid
        else:
            left = mid
    return right


n, m, q = map(int, input().split())
p_sum = [0] * (n + 1)
m_indexes = [[0] for _ in range(m + 1)]  # m_indexes[種類][番目] 全アイテム内での通し番号
m_sum = [[0] for _ in range(m + 1)]  # m_sum[種類][番目] 当該種類番号での累積和

for i in range(n):
    p, a = map(int, input().split())
    p_sum[i + 1] = p_sum[i] + p
    m_indexes[a].append(i + 1)
    m_sum[a].append(m_sum[a][-1] + p)

for _ in range(q):
    t, l, r = map(int, input().split())
    ans = p_sum[r] - p_sum[l - 1]

    left_bound = lower_bound(m_indexes[t], l)
    right_bound = lower_bound(m_indexes[t], r)

    if right_bound == len(m_sum[t]):
        ans -= (m_sum[t][right_bound - 1] - m_sum[t][left_bound - 1]) // 2
    elif r < m_indexes[t][right_bound]:
        ans -= (m_sum[t][right_bound - 1] - m_sum[t][left_bound - 1]) // 2
    else:
        ans -= (m_sum[t][right_bound] - m_sum[t][left_bound - 1]) // 2
    print(ans)

このプログラムの計算量はO(N+Q log N)です。小課題をすべて合格できるので100点が得られます。