AtCoderのABC    次のABCの問題へ    前のABCの問題へ

ABC267-E Erasing Vertices 2


問題へのリンク


C#のソース(AVL木を使用)

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("4 3");
            WillReturn.Add("3 1 4 2");
            WillReturn.Add("1 2");
            WillReturn.Add("1 3");
            WillReturn.Add("4 1");
            //3
        }
        else if (InputPattern == "Input2") {
            WillReturn.Add("7 13");
            WillReturn.Add("464 661 847 514 74 200 188");
            WillReturn.Add("5 1");
            WillReturn.Add("7 1");
            WillReturn.Add("5 7");
            WillReturn.Add("4 1");
            WillReturn.Add("4 5");
            WillReturn.Add("2 4");
            WillReturn.Add("5 2");
            WillReturn.Add("1 3");
            WillReturn.Add("1 6");
            WillReturn.Add("3 5");
            WillReturn.Add("1 2");
            WillReturn.Add("4 6");
            WillReturn.Add("2 7");
            //1199
        }
        else {
            string wkStr;
            while ((wkStr = Console.ReadLine()) != null) WillReturn.Add(wkStr);
        }
        return WillReturn;
    }

    // 隣接リスト
    static Dictionary<long, List<long>> mToNodeListDict = new Dictionary<long, List<long>>();

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

        long[] wkArr = { };
        Action<string> SplitAct = pStr =>
            wkArr = pStr.Split(' ').Select(pX => long.Parse(pX)).ToArray();

        SplitAct(InputList[0]);
        long N = wkArr[0];

        foreach (string EachStr in InputList.Skip(2)) {
            SplitAct(EachStr);
            long FromNode = wkArr[0];
            long ToNode = wkArr[1];

            if (mToNodeListDict.ContainsKey(FromNode) == false) {
                mToNodeListDict[FromNode] = new List<long>();
            }
            if (mToNodeListDict.ContainsKey(ToNode) == false) {
                mToNodeListDict[ToNode] = new List<long>();
            }
            mToNodeListDict[FromNode].Add(ToNode);
            mToNodeListDict[ToNode].Add(FromNode);
        }

        // コスト[ノード]なDict
        var CostDict = new Dictionary<long, long>();
        for (long I = 1; I <= N; I++) {
            CostDict[I] = 0;
        }

        foreach (var EachPair in mToNodeListDict) {
            foreach (long EachToNode in EachPair.Value) {
                CostDict[EachPair.Key] += AArr[EachToNode - 1];
            }
        }

        // AVL木で{ノード , コスト } を管理 (コストがソートキー)
        var Ins_AVL_Set_MultiSet = new AVL_Set_MultiSet<SortableStruct>();
        Ins_AVL_Set_MultiSet.IsMultiSet = true;

        foreach (var EachPair in CostDict) {
            SortableStruct WillAdd;
            WillAdd.Node = EachPair.Key;
            WillAdd.Cost = EachPair.Value;
            Ins_AVL_Set_MultiSet.Add(WillAdd);
        }

        var CostList = new List<long>();
        while (Ins_AVL_Set_MultiSet.Count > 0) {
            SortableStruct FirstItem = Ins_AVL_Set_MultiSet[0];

            long CurrNode = FirstItem.Node;
            CostList.Add(FirstItem.Cost);
            Ins_AVL_Set_MultiSet.RemoveAt(0);
            CostDict.Remove(CurrNode);

            if (mToNodeListDict.ContainsKey(CurrNode)) {
                foreach (long EachNearNode in mToNodeListDict[CurrNode]) {
                    if (CostDict.ContainsKey(EachNearNode)) {
                        SortableStruct BeforeStruct;
                        BeforeStruct.Cost = CostDict[EachNearNode];
                        BeforeStruct.Node = EachNearNode;

                        CostDict[EachNearNode] -= AArr[CurrNode - 1];
                        SortableStruct AfterStruct;
                        AfterStruct.Cost = CostDict[EachNearNode];
                        AfterStruct.Node = EachNearNode;

                        Ins_AVL_Set_MultiSet.Remove(BeforeStruct);
                        Ins_AVL_Set_MultiSet.Add(AfterStruct);
                    }
                }
            }
        }
        Console.WriteLine(CostList.Max());
    }
}

////////////////////////////////////////////////////////////////
// ソートを定義した構造体
////////////////////////////////////////////////////////////////
struct SortableStruct : IComparable<SortableStruct>
{
    internal long Cost;
    internal long Node;

    // OrderBy Cost asc でソート
    public int CompareTo(SortableStruct pOtherIns)
    {
        return Cost.CompareTo(pOtherIns.Cost);
    }
}

#region AVL_Set_MultiSet
/// <summary>
/// 要素の追加、削除、検索、取得が可能な集合を表します.
/// </summary>
/// <typeparam name="T">優先度付きキュー内の要素の型を指定します.</typeparam>
/// <remarks>内部的にはAVL木によって実装されています.</remarks>
internal class AVL_Set_MultiSet<T>
{
    Node root;
    readonly IComparer<T> comparer;
    readonly Node nil;

    /// <summary>
    /// 多重集合かどうかを表します.
    /// </summary>
    internal bool IsMultiSet { get; set; }
    internal AVL_Set_MultiSet(IComparer<T> comparer)
    {
        nil = new Node(default(T));
        root = nil;
        this.comparer = comparer;
    }

    internal AVL_Set_MultiSet() : this(Comparer<T>.Default) { }

    /// <summary>
    /// 要素をコレクションに追加します.
    /// </summary>
    /// <remarks>この操作は計算量 O(log N) で実行されます.</remarks>
    internal bool Add(T v)
    {
        return insert(ref root, v);
    }

    /// <summary>
    /// v が存在するならコレクションから削除します.
    /// </summary>
    /// <remarks>この操作は計算量 O(log N) で実行されます.</remarks>
    internal bool Remove(T v)
    {
        return remove(ref root, v);
    }

    /// <summary>
    /// 0-indexed で index 番目の要素をコレクションから取得します..
    /// </summary>
    /// <remarks>この操作は計算量 O(log N) で実行されます.</remarks>
    internal T this[int index] { get { return find(root, index); } }
    internal int Count { get { return root.Count; } }

    internal void RemoveAt(int k)
    {
        if (k < 0 || k >= root.Count) throw new ArgumentOutOfRangeException();
        removeAt(ref root, k);
    }

    /// <summary>
    /// このコレクションに含まれる要素を昇順に並べて返します.
    /// </summary>
    /// <remarks>この操作は計算量 O(N) で実行されます.</remarks>
    internal T[] Items
    {
        get
        {
            T[] ret = new T[root.Count];
            int k = 0;
            walk(root, ret, ref k);
            return ret;
        }
    }

    private void walk(Node t, T[] a, ref int k)
    {
        if (t.Count == 0) return;
        walk(t.lst, a, ref k);
        a[k++] = t.Key;
        walk(t.rst, a, ref k);
    }

    private bool insert(ref Node t, T key)
    {
        if (t.Count == 0) { t = new Node(key); t.lst = t.rst = nil; t.Update(); return true; }
        int cmp = comparer.Compare(t.Key, key);
        bool res;
        if (cmp > 0)
            res = insert(ref t.lst, key);
        else if (cmp == 0) {
            if (IsMultiSet) res = insert(ref t.lst, key);
            else return false;
        }
        else res = insert(ref t.rst, key);
        balance(ref t);
        return res;
    }

    private bool remove(ref Node t, T key)
    {
        if (t.Count == 0) return false;
        int cmp = comparer.Compare(key, t.Key);
        bool ret;
        if (cmp < 0) ret = remove(ref t.lst, key);
        else if (cmp > 0) ret = remove(ref t.rst, key);
        else {
            ret = true;
            var k = t.lst.Count;
            if (k == 0) { t = t.rst; return true; }
            if (t.rst.Count == 0) { t = t.lst; return true; }

            t.Key = find(t.lst, k - 1);
            removeAt(ref t.lst, k - 1);
        }
        balance(ref t);
        return ret;
    }

    private void removeAt(ref Node t, int k)
    {
        int cnt = t.lst.Count;
        if (cnt < k) removeAt(ref t.rst, k - cnt - 1);
        else if (cnt > k) removeAt(ref t.lst, k);
        else {
            if (cnt == 0) { t = t.rst; return; }
            if (t.rst.Count == 0) { t = t.lst; return; }

            t.Key = find(t.lst, k - 1);
            removeAt(ref t.lst, k - 1);
        }
        balance(ref t);
    }

    private void balance(ref Node t)
    {
        int balance = t.lst.Height - t.rst.Height;
        if (balance == -2) {
            if (t.rst.lst.Height - t.rst.rst.Height > 0) { rotR(ref t.rst); }
            rotL(ref t);
        }
        else if (balance == 2) {
            if (t.lst.lst.Height - t.lst.rst.Height < 0) rotL(ref t.lst);
            rotR(ref t);
        }
        else t.Update();
    }

    private T find(Node t, int k)
    {
        if (k < 0 || k > root.Count) throw new ArgumentOutOfRangeException();
        while (true) {
            if (k == t.lst.Count) return t.Key;
            else if (k < t.lst.Count) t = t.lst;
            else { k -= t.lst.Count + 1; t = t.rst; }
        }
    }
    /// <summary>
    /// コレクションに含まれる要素であって、 v 以上の最小の要素の番号を返します。
    /// </summary>
    /// <remarks>この操作は計算量 O(log N) で実行されます.</remarks>
    internal int LowerBound(T v)
    {
        // 追加機能 件数が0なら-1を返す
        if (this.Count == 0) {
            return -1;
        }

        // 追加機能 v 以上な要素が無い場合は-1を返す
        T MaxVal = this[this.Count - 1];
        int cmp = comparer.Compare(MaxVal, v);
        if (cmp == -1) return -1;

        int k = 0;
        Node t = root;
        while (true) {
            if (t.Count == 0) return k;
            if (comparer.Compare(v, t.Key) <= 0) t = t.lst;
            else { k += t.lst.Count + 1; t = t.rst; }
        }
    }
    /// <summary>
    /// コレクションに含まれる要素であって、 v より真に大きい、最小の要素の番号を返します。
    /// </summary>
    /// <remarks>この操作は計算量 O(log N) で実行されます.</remarks>
    internal int UpperBound(T v)
    {
        // 追加機能 件数が0なら-1を返す
        if (this.Count == 0) {
            return -1;
        }

        // 追加機能 v 超えな要素が無い場合は-1を返す
        T MaxVal = this[this.Count - 1];
        int cmp = comparer.Compare(MaxVal, v);
        if (cmp <= 0) return -1;

        int k = 0;
        Node t = root;
        while (true) {
            if (t.Count == 0) return k;
            if (comparer.Compare(t.Key, v) <= 0) { k += t.lst.Count + 1; t = t.rst; }
            else t = t.lst;
        }
    }

    // 追加機能 V未満で最大の要素の番号を返す
    internal int Lower_Max(T v)
    {
        // 件数が0なら-1を返す
        if (this.Count == 0) {
            return -1;
        }

        // v 未満な要素が無い場合は-1を返す
        T MinVal = this[0];
        int cmp = comparer.Compare(MinVal, v);
        if (cmp >= 0) return -1;

        // v 以上の件数を調べる
        int More_Or_Equal_cnt = this.Count;
        int UB = this.Count - 1;
        int LowerB = LowerBound(v);
        if (IsValidInd(LowerB)) {
            More_Or_Equal_cnt -= (UB - LowerB + 1);
        }
        return More_Or_Equal_cnt - 1;
    }

    // 追加機能 V以下で最大の要素の番号を返す
    internal int LowerOrEqual_Max(T v)
    {
        // 件数が0なら-1を返す
        if (this.Count == 0) {
            return -1;
        }

        // v 以下な要素が無い場合は-1を返す
        T MinVal = this[0];
        int cmp = comparer.Compare(MinVal, v);
        if (cmp > 0) return -1;

        // v 超えの件数を調べる
        int More_cnt = this.Count;
        int UB = this.Count - 1;
        int UpperB = UpperBound(v);
        if (IsValidInd(UpperB)) {
            More_cnt -= (UB - UpperB + 1);
        }
        return More_cnt - 1;
    }

    // 追加機能 LowerBoundなどで返したIndが、有効範囲かを判定
    internal bool IsValidInd(int pInd)
    {
        if (pInd < 0) return false;
        if (this.Count <= pInd) return false;
        return true;
    }

    private void rotR(ref Node t)
    {
        Node l = t.lst;
        t.lst = l.rst;
        l.rst = t;
        t.Update();
        l.Update();
        t = l;
    }

    private void rotL(ref Node t)
    {
        Node r = t.rst;
        t.rst = r.lst;
        r.lst = t;
        t.Update();
        r.Update();
        t = r;
    }

    class Node
    {
        internal Node(T key)
        {
            Key = key;
        }
        internal int Count { get; private set; }
        internal int Height { get; private set; }
        internal T Key { get; set; }
        internal Node lst, rst;
        internal void Update()
        {
            Count = 1 + lst.Count + rst.Count;
            Height = 1 + Math.Max(lst.Height, rst.Height);
        }
        public override string ToString()
        {
            return string.Format("Count = {0}, Key = {1}", Count, Key);
        }
    }
}
#endregion


C#のソース(セグメント木を使用)

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

// ABC267-E Erasing Vertices 2
// https://atcoder.jp/contests/abc267/tasks/abc267_e
class Program
{
    static string InputPattern = "InputX";

    static List<string> GetInputList()
    {
        var WillReturn = new List<string>();
        string wkStr;
        while ((wkStr = Console.ReadLine()) != null) WillReturn.Add(wkStr);
        return WillReturn;
    }

    // 隣接リスト
    static Dictionary<long, List<long>> mToNodeListDict = new Dictionary<long, List<long>>();

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

        long[] wkArr = { };
        Action<string> SplitAct = pStr =>
            wkArr = pStr.Split(' ').Select(pX => long.Parse(pX)).ToArray();

        SplitAct(InputList[0]);
        long N = wkArr[0];

        foreach (string EachStr in InputList.Skip(2)) {
            SplitAct(EachStr);
            long FromNode = wkArr[0];
            long ToNode = wkArr[1];

            if (mToNodeListDict.ContainsKey(FromNode) == false) {
                mToNodeListDict[FromNode] = new List<long>();
            }
            if (mToNodeListDict.ContainsKey(ToNode) == false) {
                mToNodeListDict[ToNode] = new List<long>();
            }
            mToNodeListDict[FromNode].Add(ToNode);
            mToNodeListDict[ToNode].Add(FromNode);
        }

        // コスト[ノード]なDict
        var CostDict = new Dictionary<long, long>();
        for (long I = 1; I <= N; I++) {
            CostDict[I] = 0;
        }

        foreach (var EachPair in mToNodeListDict) {
            foreach (long EachToNode in EachPair.Value) {
                CostDict[EachPair.Key] += AArr[EachToNode - 1];
            }
        }

        var InsSegmentTree = new SegmentTree(N, long.MaxValue);
        for (long I = 1; I <= N; I++) {
            InsSegmentTree.Update(I, CostDict[I]);
        }

        var CostList = new List<long>();
        for (long I = 1; I <= N; I++) {
            long MinVal, MinValInd;
            InsSegmentTree.GetMinInfo(out MinVal, out MinValInd);
            CostList.Add(MinVal);
            InsSegmentTree.Update(MinValInd, long.MaxValue);

            if (mToNodeListDict.ContainsKey(MinValInd)) {
                foreach (long EachNearNode in mToNodeListDict[MinValInd]) {
                    long CurrVal = InsSegmentTree.Internal_Query(EachNearNode, EachNearNode);
                    InsSegmentTree.Update(EachNearNode, CurrVal - AArr[MinValInd - 1]);
                }
            }
        }
        Console.WriteLine(CostList.Max());
    }
}

#region SegmentTree
// SegmentTreeクラス (RMQ and 1点更新)
internal class SegmentTree
{
    // 拡張機能 (最小値と、最小値を持つIndを返す)
    // 最小値が複数あったら、左のIndを優先
    internal void GetMinInfo(out long pMinVal, out long pMinValInd)
    {
        // 全区間での最小値を求める
        pMinVal = Internal_Query(0, mExternalArrUB);

        // 二分探索を行う
        long L = 0, R = mExternalArrUB;

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

            if (LeftMin == pMinVal) R = Mid;
            else L = Mid;
        }
        pMinValInd = ((Internal_Query(L, L) == pMinVal) ? L : R);
    }

    // 拡張機能 (最小値と、最小値を持つIndを返す)
    // 最小値が複数あったら、右のIndを優先
    internal void GetMinInfo_Right(out long pMinVal, out long pMinValInd)
    {
        // 全区間での最小値を求める
        pMinVal = Internal_Query(0, mExternalArrUB);

        // 二分探索を行う
        long L = 0, R = mExternalArrUB;

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

            if (RightMin == pMinVal) L = Mid;
            else R = Mid;
        }
        pMinValInd = ((Internal_Query(R, R) == pMinVal) ? R : L);
    }

    private long[] mTreeNodeArr;
    private long UB; // 木のノードの配列のUB
    private long mLeafCnt; // 葉ノードの数
    private long mExternalArrUB;

    // ノードの添字を引数とし、範囲の開始添字と終了添字を持つ配列
    private struct RangeInfoDef
    {
        internal long StaInd;
        internal long EndInd;
    }
    private RangeInfoDef[] mRangeInfo;

    // ノードのIndexの列挙を返す
    internal IEnumerable<long> GetNodeIndEnum()
    {
        for (long I = 0; I <= mExternalArrUB; I++) {
            yield return I;
        }
    }

    // 木のノードのUBを返す
    internal long GetUB()
    {
        return mExternalArrUB;
    }

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

        // 簡単のため、葉ノード数を2のべき乗に
        long ArrLength = 0;
        for (long I = 1; I < int.MaxValue; I *= 2) {
            ArrLength += I;
            mLeafCnt = I;

            if (pExternalArrUB + 1 < mLeafCnt) break;
        }

        // すべての値をpInitValに
        UB = ArrLength - 1;
        mTreeNodeArr = new long[UB + 1];
        for (long I = 0; I <= UB; I++) {
            mTreeNodeArr[I] = pInitVal;
        }

        // ノードの添字を引数とし、範囲の開始添字と終了添字を持つ配列の作成
        mRangeInfo = new RangeInfoDef[UB + 1];
        for (long I = 0; I <= UB; I++) {
            if (I == 0) {
                RangeInfoDef WillSet1;
                WillSet1.StaInd = 0;
                WillSet1.EndInd = mLeafCnt - 1;
                mRangeInfo[I] = WillSet1;
                continue;
            }
            long ParentNode = DeriveParentNode(I);
            RangeInfoDef ParentRangeInfo = mRangeInfo[ParentNode];

            RangeInfoDef WillSet2;
            long Mid = (ParentRangeInfo.StaInd + ParentRangeInfo.EndInd) / 2;

            if (I % 2 == 1) { // 奇数ノードの場合
                WillSet2.StaInd = ParentRangeInfo.StaInd;
                WillSet2.EndInd = Mid;
            }
            else { // 偶数ノードの場合
                WillSet2.StaInd = Mid + 1;
                WillSet2.EndInd = ParentRangeInfo.EndInd;
            }
            mRangeInfo[I] = WillSet2;
        }
    }

    // 親ノードの添字を取得
    private long DeriveParentNode(long pTarget)
    {
        return (pTarget - 1) / 2;
    }

    // 子ノードの添字(小さいほう)を取得
    private long DeriveChildNode(long pTarget)
    {
        return pTarget * 2 + 1;
    }

    // 葉ノードの配列の添字を木の添字に変換して返す
    private long DeriveTreeNode(long pLeafArrInd)
    {
        long BaseInd = UB - mLeafCnt + 1;
        return BaseInd + pLeafArrInd;
    }

    // 葉ノードの配列のK番目の値をNewValに変更
    internal void Update(long pK, long pNewVal)
    {
        long CurrNode = DeriveTreeNode(pK);
        mTreeNodeArr[CurrNode] = pNewVal;

        // 登りながら更新
        while (CurrNode > 0) {
            CurrNode = DeriveParentNode(CurrNode);
            long ChildNode1 = DeriveChildNode(CurrNode);
            long ChildNode2 = ChildNode1 + 1;
            mTreeNodeArr[CurrNode] =
                Math.Min(mTreeNodeArr[ChildNode1], mTreeNodeArr[ChildNode2]);
        }
    }

    // 開始添字と終了添字とカレントノードを引数として、最小値を返す
    internal long Internal_Query(long pSearchStaInd, long pSearchEndInd)
    {
        return Private_Query(pSearchStaInd, pSearchEndInd, 0);
    }
    private long Private_Query(long pSearchStaInd, long pSearchEndInd, long pCurrNode)
    {
        long CurrNodeStaInd = mRangeInfo[pCurrNode].StaInd;
        long CurrNodeEndInd = mRangeInfo[pCurrNode].EndInd;

        // OverLapしてなければ、long.MaxValue
        if (CurrNodeEndInd < pSearchStaInd || pSearchEndInd < CurrNodeStaInd)
            return long.MaxValue;

        // 完全に含んでいれば、このノードの値
        if (pSearchStaInd <= CurrNodeStaInd && CurrNodeEndInd <= pSearchEndInd)
            return mTreeNodeArr[pCurrNode];

        // そうでなければ、2つの子の最小値
        long ChildNode1 = DeriveChildNode(pCurrNode);
        long ChildNode2 = ChildNode1 + 1;

        long ChildVal1 = Private_Query(pSearchStaInd, pSearchEndInd, ChildNode1);
        long ChildVal2 = Private_Query(pSearchStaInd, pSearchEndInd, ChildNode2);
        return Math.Min(ChildVal1, ChildVal2);
    }

    internal void DebugPrint()
    {
        Console.WriteLine("■■■■■■■■■■");
        for (long I = 0; I <= UB; I++) {
            Console.WriteLine("mTreeNodeArr[{0}] = {1}", I, mTreeNodeArr[I]);
        }
    }
}
#endregion


解説

コスト[ノード]なDictと
ノードとコストの構造体を要素に持つAVL木を使って
最小コストから貪欲にノードを削除してます。

RMQ and 1点更新 のセグメント木で
拡張機能 (最小値と、最小値を持つIndを返す)
を使って解くこともできます。