ABC149-E Handshake

問題へのリンク

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("5 3");
            WillReturn.Add("10 14 19 34 33");
            //202
        }
        else if (InputPattern == "Input2") {
            WillReturn.Add("9 14");
            WillReturn.Add("1 3 5 110 24 21 34 5 3");
            //1837
        }
        else if (InputPattern == "Input3") {
            WillReturn.Add("9 73");
            WillReturn.Add("67597 52981 5828 66249 75177 64141 40773 79105 16076");
            //8128170
        }
        else {
            string wkStr;
            while ((wkStr = Console.ReadLine()) != null) WillReturn.Add(wkStr);
        }
        return WillReturn;
    }

    static long mM;

    static long[] mAArr;
    static long[] mRunSumArr;

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

        mAArr = InputList[1].Split(' ').Select(pX => long.Parse(pX)).ToArray();
        Array.Sort(mAArr);

        // 二分法で降順でのM番目の値を求める
        long L = 0;
        long R = mAArr.Max() * 2 + 1;

        while (L + 1 < R) {
            long Mid = (L + R) / 2;

            if (IsUpperOrEqual(Mid)) {
                L = Mid;
            }
            else {
                R = Mid;
            }
        }
        long Shikii = L;

        // 累積和を求める
        mRunSumArr = (long[])mAArr.Clone();
        for (long I = 1; I <= mRunSumArr.GetUpperBound(0); I++) {
            mRunSumArr[I] += mRunSumArr[I - 1];
        }

        // 指定した値超えの件数と、和を求める
        long OverCnt;
        long OverSum;
        DeriveCntAndSum(Shikii, out OverCnt, out OverSum);

        long Answer = OverSum;
        long RestCnt = mM - OverCnt;
        Answer += Shikii * RestCnt;
        Console.WriteLine(Answer);
    }

    // Valを引数として、順位がM位、または、M位より悪いか判定
    static bool IsUpperOrEqual(long pVal)
    {
        // 引数と同じか、良い順位の値の件数
        long UpperOrEqualCntSum = 0;

        foreach (long EachA in mAArr) {
            long CurrLowerOrEqualCnt = DeriveUpperOrEqualCnt(pVal, EachA);
            UpperOrEqualCntSum += CurrLowerOrEqualCnt;
            if (UpperOrEqualCntSum >= mM) return true;
        }
        return UpperOrEqualCntSum >= mM;
    }

    // ValとAArrの要素を引数として、和がVal以上になる件数を返す
    static long DeriveUpperOrEqualCnt(long pVal, long pA)
    {
        // 初項がpVal以上の特殊ケース
        if (mAArr[0] + pA >= pVal) {
            return mAArr.Length;
        }
        // 末項がpVal未満の特殊ケース
        if (mAArr.Last() + pA < pVal) {
            return 0;
        }

        long L = 0;
        long R = mAArr.GetUpperBound(0);

        while (L + 1 < R) {
            long Mid = (L + R) / 2;
            long MidVal = mAArr[Mid] + pA;

            if (MidVal < pVal) {
                L = Mid;
            }
            else {
                R = Mid;
            }
        }
        return mAArr.GetUpperBound(0) - R + 1;
    }

    // 閾値超えの、件数と和を求める
    static void DeriveCntAndSum(long pShikii, out long pOverCnt, out long pOverSum)
    {
        pOverCnt = 0;
        pOverSum = 0;

        long AllSum = mAArr.Sum();
        long UB = mAArr.GetUpperBound(0);

        foreach (long EachA in mAArr) {
            long MinInd = DeriveMinInd(pShikii, EachA);
            if (MinInd == -1) continue;

            long CurrCnt = UB - MinInd + 1;
            long CurrSum = AllSum;
            if (MinInd > 0) {
                CurrSum -= mRunSumArr[MinInd - 1];
            }
            CurrSum += EachA * CurrCnt;

            pOverCnt += CurrCnt;
            pOverSum += CurrSum;
        }
    }

    // 閾値とAArrの要素を引数として、和が閾値超えになる最小の添字を返す
    static long DeriveMinInd(long pShikii, long pA)
    {
        // 初項が閾値超えの特殊ケース
        if (mAArr[0] + pA > pShikii) {
            return 0;
        }
        // 末項が閾値以下の特殊ケース
        if (mAArr.Last() + pA <= pShikii) {
            return -1;
        }

        long L = 0;
        long R = mAArr.GetUpperBound(0);

        while (L + 1 < R) {
            long Mid = (L + R) / 2;
            long MidVal = mAArr[Mid] + pA;

            if (MidVal <= pShikii) {
                L = Mid;
            }
            else {
                R = Mid;
            }
        }
        return R;
    }
}

解説

最初に、二分法で
降順でのM番目の値を求めます。
この値を閾値とします。

次に、二分法で
閾値超えの値の、件数と和
を求めます。

これらの情報が分かれば、解も求めることができます。