ABC351-F Double Sum

問題へのリンク

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");
            WillReturn.Add("2 5 3");
            //4
        }
        else if (InputPattern == "Input2") {
            WillReturn.Add("10");
            WillReturn.Add("5 9 3 0 4 8 7 5 4 0");
            //58
        }
        else {
            string wkStr;
            while ((wkStr = Console.ReadLine()) != null) WillReturn.Add(wkStr);
        }
        return WillReturn;
    }

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

        // 座圧後の値[座圧前の値]
        Dictionary<long, long> ZaatuDict = DeriveZaatuDict(AArr);

        long Fenwick_Tree_UB = ZaatuDict.Count - 1;

        var Ins_Fenwick_Tree1 = new Fenwick_Tree(Fenwick_Tree_UB);
        var Ins_Fenwick_Tree2 = new Fenwick_Tree(Fenwick_Tree_UB);
        foreach (long EachA in AArr) {
            long Ind = ZaatuDict[EachA];
            Ins_Fenwick_Tree1.Add(Ind, EachA);
            Ins_Fenwick_Tree2.Add(Ind, 1);
        }

        long Answer = 0;
        for (long I = 0; I <= AArr.GetUpperBound(0) - 1; I++) {
            long RangeSta = ZaatuDict[AArr[I]] + 1;
            long RangeEnd = Fenwick_Tree_UB;

            if (RangeSta <= RangeEnd) {
                long CurrSum = 0;
                long CurrCnt = 0;

                CurrSum = Ins_Fenwick_Tree1.GetSum(RangeSta, RangeEnd);
                CurrCnt = Ins_Fenwick_Tree2.GetSum(RangeSta, RangeEnd);
                CurrSum -= AArr[I] * CurrCnt;

                Answer += CurrSum;
            }

            long Ind = ZaatuDict[AArr[I]];
            Ins_Fenwick_Tree1.Add(Ind, -AArr[I]);
            Ins_Fenwick_Tree2.Add(Ind, -1);
        }
        Console.WriteLine(Answer);
    }

    //////////////////////////////////////////////////////////////////////////
    // 列挙を引数として、座標圧縮し、座圧後の値[座圧前の値]なDictを返す
    //////////////////////////////////////////////////////////////////////////
    static Dictionary<long, long> DeriveZaatuDict(IEnumerable<long> pEnum)
    {
        var ZaatuDict = new Dictionary<long, long>();
        var ValSet = new HashSet<long>(pEnum);
        long No = 0;
        foreach (long EachVal in ValSet.OrderBy(pX => pX)) {
            ZaatuDict[EachVal] = No;
            No++;
        }
        return ZaatuDict;
    }
}

// フェニック木
#region Fenwick_Tree
internal class Fenwick_Tree
{
    private long[] mBitArr;
    private long mExternalArrUB;

    // コンストラクタ
    internal Fenwick_Tree(long pExternalArrUB)
    {
        mExternalArrUB = pExternalArrUB;

        // フェニック木の外部配列は0オリジンで、
        // フェニック木の内部配列は1オリジンなため、2を足す
        mBitArr = new long[pExternalArrUB + 2];
    }

    // [pSta,pEnd] のSumを返す
    internal long GetSum(long pSta, long pEnd)
    {
        return GetSum(pEnd) - GetSum(pSta - 1);
    }

    // [0,pEnd] のSumを返す
    internal long GetSum(long pEnd)
    {
        pEnd++; // 1オリジンに変更

        long Sum = 0;
        while (pEnd >= 1) {
            Sum += mBitArr[pEnd];
            pEnd -= pEnd & -pEnd;
        }
        return Sum;
    }

    // [I] に Xを加算
    internal void Add(long pI, long pX)
    {
        pI++; // 1オリジンに変更

        while (pI <= mBitArr.GetUpperBound(0)) {
            mBitArr[pI] += pX;
            pI += pI & -pI;
        }
    }
}
#endregion

解説

添字は座標圧縮し、座標圧縮前の数の総和なフェニック木
添字は座標圧縮し、座標圧縮前の数の個数なフェニック木
で2つのフェニック木を用意して解いてます。