AtCoder Regular Contest 116 B – Products of Min-Max をPython3で解く

Share

AtCoder上にある問題のうち、AtCoder Problemsでdiff 800以上と判定されているものを順番に解いていく企画。
基本的な考え方は全てコード中のコメントに入れてあるので、参照のこと。

出典:
AtCoder Regular Contest 116 B – Products of Min-Max


Aの空でない部分列Bは\(2^N – 1\)個ありますとか見た瞬間に、ヤバい香りしかしない問題ではある。
数学力が高い人は公式解説で理解できるのだと思うが、僕のように数学力が0の人は、Σ記号を見るだけで頭痛がすることだろう。

というわけで、かなり噛み砕いた考え方をコメントに入れてあるので、参考にしてみてもらいたい。

# AtCoder Regular Contest 116 B - Products of Min-Max
# https://atcoder.jp/contests/arc116/tasks/arc116_b
# tag: 考察 数列 部分列 最大値 最小値 順列・組み合わせ 総和 MOD

# 1 <= N <= 2 * 10^5 の制約が厳しく、正直にそのまま実装するのは
# もちろんのこと、(L, R) を決めるような O(N^2) のやり方でも
# 時間切れになりそう。

# まず、そのままだと考察しづらいので A をソート済みとする。
# 仮に A = (1, 2, 3, 4, 5) のケースで考えてみると、
# 最小値 1 を固定した時、うまく計算できないだろうか?

# 合計されるのは、1 ~ それぞれの最大値の間に挟まる数が
# それぞれある場合、ない場合の 2 通りずつになるので、
# 最大値 1: 1 * 1 - 1回 (数字が 1 しか残っていないケース)
# 最大値 2: 1 * 2 - 1回 (数字が 1, 2 しか残っていないケース)
# 最大値 3: 1 * 3 - 2回 (1 ~ 3 の間の 2 が 残る/残らない ケース)
# 最大値 4: 1 * 4 - 2^2回 (1 ~ 4 の間の 2, 3 がそれぞれ 残る/残らない)
# 最大値 5: 1 * 5 - 2^3回 (1 ~ 5 の間の 2, 3, 4 が以下略)
# となる。

# 同様に最小値が 2 の場合を考えると、
# 最大値 2: 2 * 2 - 1回
# 最大値 3: 2 * 3 - 1回
# 最大値 4: 2 * 4 - 2回
# 最大値 5: 2 * 5 - 2^2回

# 数字が一つしか残らない場合を別に扱うと、
# 最小値 1 の場合:
# 1 * 1 + 1 * (2 * 1 + 3 * 2 + 4 * 2^2 + 5 * 2^3)
# 最小値 2 の場合:
# 2 * 2 + 2 * (3 * 1 + 4 * 2 + 5 * 2^2)
# 以下同様に、
# 3 * 3 + 3 * (4 * 1 + 5 * 2)
# 4 * 4 + 4 * (5 * 1)
# 5 * 5 + 5 * (0)

# ここで一番右のカッコ内に注目すると、
# 最小値が減るに従って、2 倍して前の最大値を足した値になっている
# これは O(1) で求められるので、全体でも O(N) でいけそう。

def main():
    N = int(input())
    A = list(map(int, input().split()))
    MOD = 998244353

    # まずは数列をソート
    A.sort()

    result = 0

    # 先に数字が一つしか残らないケースを加算しておく
    for i in range(N):
        result = (result + A[i] ** 2) % MOD

    # 最小値を右から 2 番目~最左端へと動かしつつ計算していく
    # 上記の右カッコ内を管理する変数を用意
    mult = 0
    for min_i in range(N-2, -1, -1):
        # 右カッコ内を更新
        mult = (mult * 2 + A[min_i+1]) % MOD
        # 値を計算して足していく
        result = (result + A[min_i] * mult) % MOD
    
    print(result)

main()
Share

コメントを残す

メールアドレスが公開されることはありません。