ABC216-E Amusement Park

問題へのリンク

C#のソース

using System;
using System.Collections.Generic;
using System.Linq;

class Program
{
    static string InputPattern = "InputX";

    static List<string> GetInputList()
    {
        var WillReturn = new List<string>();

        if (InputPattern == "Input1") {
            WillReturn.Add("3 5");
            WillReturn.Add("100 50 102");
            //502
        }
        else if (InputPattern == "Input2") {
            WillReturn.Add("2 2021");
            WillReturn.Add("2 3");
            //9
        }
        else {
            string wkStr;
            while ((wkStr = Console.ReadLine()) != null) WillReturn.Add(wkStr);
        }
        return WillReturn;
    }

    static long mK;
    static long[] mAArr;
    static long UB;

    static void Main()
    {
        List<string> InputList = GetInputList();
        long[] wkArr = InputList[0].Split(' ').Select(pX => long.Parse(pX)).ToArray();
        mK = wkArr[1];

        mAArr = InputList[1].Split(' ').Select(pX => long.Parse(pX)).ToArray();
        UB = mAArr.GetUpperBound(0);

        // 昇順にソートしておく
        Array.Sort(mAArr);
        UB = mAArr.GetUpperBound(0);

        long Shikii = ExecNibunhou(mAArr);
        //Console.WriteLine("閾値={0}", Shikii);

        long Answer = 0;
        long Rest = mK;
        for (long I = UB; 0 <= I; I--) {
            if (mAArr[I] < Shikii) break;

            // 等差数列の和の公式
            long Kousuu = mAArr[I] - Shikii + 1;
            Answer += (Shikii + mAArr[I]) * Kousuu / 2;
            Rest -= Kousuu;
        }
        //Console.WriteLine("閾値以上のSum={0}", Answer);

        // 残りの取得分
        if (Shikii > 0) {
            Answer += Rest * (Shikii - 1);
        }

        Console.WriteLine(Answer);
    }

    // 二分法で、閾値以上のマスを全て埋めれる、最小の閾値を返す
    static long ExecNibunhou(long[] pArr)
    {
        // 閾値1を 達成可能な特殊ケース
        if (DeriveOverMasuCnt(1) <= mK) {
            return 1;
        }

        long L = 1;
        long R = pArr.Last() + 1;

        while (L + 1 < R) {
            long Mid = (L + R) / 2;
            if (DeriveOverMasuCnt(Mid) <= mK) {
                R = Mid;
            }
            else {
                L = Mid;
            }
        }
        return R;
    }

    // 引数の閾値以上のマスの数を返す
    static long DeriveOverMasuCnt(long pVal)
    {
        long WillReturn = 0;

        for (long I = UB; 0 <= I; I--) {
            if (mAArr[I] < pVal) break;

            WillReturn += mAArr[I] - pVal + 1;
        }
        return WillReturn;
    }
}

解説

昇順にソートしても問題ないので
昇順にソートして棒グラフで考えます。

　　　□□
　　　□□
　　□□□
　□□□□
□□□□□

閾値を決めて、閾値以上のマスがK個以下かは、
単調性があるので、二分探索で最小の閾値を求めます。

閾値が決まったら、閾値以上の合計は、
等差数列の和の公式で求まって、
残りの取得分の和も計算で求まります。