ABC239-F Construct Highway

問題へのリンク

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

    static void Main()
    {
        List<string> InputList = GetInputList();

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

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

        int[] DArr = InputList[1].Split(' ').Select(pX => int.Parse(pX)).ToArray();

        // 残り次数[ノード]なDict
        var DDict = new Dictionary<int, int>();
        for (int I = 0; I <= DArr.GetUpperBound(0); I++) {
            DDict[I + 1] = DArr[I];
        }

        // 木[ノード]を管理するUnionFindクラス
        var InsUnionFind = new UnionFind();
        for (int I = 1; I <= N; I++) {
            InsUnionFind.MakeSet(I);
        }

        int NeedEdge = N - 1;

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

            int Root1 = InsUnionFind.FindSet(FromNode);
            int Root2 = InsUnionFind.FindSet(ToNode);

            if (Root1 == Root2) {
                Console.WriteLine(-1);
                return;
            }

            InsUnionFind.Unite(Root1, Root2);
            DDict[FromNode]--;
            DDict[ToNode]--;

            NeedEdge--;
        }

        // 木なので枝数*2 = 次数 が必要条件
        int DegSum = DDict.Values.Sum();
        if (DegSum != NeedEdge * 2) {
            Console.WriteLine(-1);
            return;
        }

        // 次数にマイナスがあったらNG
        if (DDict.Values.Any(pX => pX < 0)) {
            Console.WriteLine(-1);
            return;
        }

        // 接続可能なノードのSet[木のID]なDict
        var CanConnSetDict = new Dictionary<int, HashSet<int>>();

        // 接続可能なノード数[木のID]なDict
        var CanConnCntDict = new Dictionary<int, int>();

        for (int I = 1; I <= N; I++) {
            int Root = InsUnionFind.FindSet(I);
            if (DDict[I] == 0) continue;

            if (CanConnSetDict.ContainsKey(Root) == false) {
                CanConnSetDict[Root] = new HashSet<int>();
            }
            CanConnSetDict[Root].Add(I);

            if (CanConnCntDict.ContainsKey(Root) == false) {
                CanConnCntDict[Root] = 0;
            }
            CanConnCntDict[Root] += DDict[I];
        }

        // 接続可能なノード数 と 木のID のペアのAVL木
        var Ins_AVL_Set_MultiSet = new AVL_Set_MultiSet<long>();
        foreach (var EachPair in CanConnCntDict) {
            int Cnt = EachPair.Value;
            int TreeID = EachPair.Key;
            Ins_AVL_Set_MultiSet.Add(GetAvlKey(Cnt, TreeID));
        }

        // 残り次数[ノード]なDict
        // var DDict = new Dictionary<int, int>();

        // 木[ノード]を管理するUnionFindクラス
        // var InsUnionFind = new UnionFind();

        // 接続可能なノードのSet[木のID]なDict
        // var CanConnSetDict = new Dictionary<int, HashSet<int>>();

        // 接続可能なノード数 と 木のID のペアのAVL木
        // var Ins_AVL_Set_MultiSet = new AVL_Set_MultiSet<long>();

        var sb = new System.Text.StringBuilder();
        while (Ins_AVL_Set_MultiSet.Count > 0) {
            if (Ins_AVL_Set_MultiSet.Count == 1) break;

            long Hash1 = Ins_AVL_Set_MultiSet[0];
            long Hash2 = Ins_AVL_Set_MultiSet[Ins_AVL_Set_MultiSet.Count - 1];
            Ins_AVL_Set_MultiSet.Remove(Hash1);
            Ins_AVL_Set_MultiSet.Remove(Hash2);

            int Tree1, Cnt1;
            GetInt2(Hash1, out Cnt1, out Tree1);
            int Tree2, Cnt2;
            GetInt2(Hash2, out Cnt2, out Tree2);

            if (CanConnSetDict[Tree1].Count == 0 || CanConnSetDict[Tree2].Count == 0) {
                break;
            }

            int ConnNode1 = CanConnSetDict[Tree1].First();
            int ConnNode2 = CanConnSetDict[Tree2].First();
            DDict[ConnNode1]--;
            DDict[ConnNode2]--;
            if (DDict[ConnNode1] == 0) {
                CanConnSetDict[Tree1].Remove(ConnNode1);
            }
            if (DDict[ConnNode2] == 0) {
                CanConnSetDict[Tree2].Remove(ConnNode2);
            }

            sb.AppendFormat("{0} {1}", ConnNode1, ConnNode2);
            sb.AppendLine();
            //Console.WriteLine("{0} {1}", ConnNode1, ConnNode2);

            InsUnionFind.Unite(Tree1, Tree2);
            Cnt1--; Cnt2--;
            int NewCnt = Cnt1 + Cnt2;
            int NewRoot = InsUnionFind.FindSet(Tree1);
            long NewHash = GetAvlKey(NewCnt, NewRoot);
            if (NewCnt > 0) {
                Ins_AVL_Set_MultiSet.Add(NewHash);
            }
            if (NewRoot == Tree1) {
                CanConnSetDict[Tree1].UnionWith(CanConnSetDict[Tree2]);
                CanConnSetDict.Remove(Tree2);
            }
            else {
                CanConnSetDict[Tree2].UnionWith(CanConnSetDict[Tree1]);
                CanConnSetDict.Remove(Tree1);
            }
        }

        if (Ins_AVL_Set_MultiSet.Count > 0) {
            Console.WriteLine(-1);
        }
        else {
            Console.Write(sb.ToString());
        }
    }

    // Int型2つからLong型のハッシュ値を求める
    static long GetAvlKey(int p1, int p2)
    {
        return (long)p1 * 1000000000 + p2;
    }

    // ハッシュからInt型2つを求める
    static void GetInt2(long pHash, out int p1, out int p2)
    {
        p1 = (int)(pHash / 1000000000);
        p2 = (int)(pHash % 1000000000);
    }
}

#region UnionFind
// UnionFindクラス
internal class UnionFind
{
    private class NodeInfoDef
    {
        internal int ParentNode;
        internal int Rank;
    }
    private Dictionary<int, NodeInfoDef> mNodeInfoDict =
        new Dictionary<int, NodeInfoDef>();

    // 要素が1つである木を森に追加
    internal void MakeSet(int pNode)
    {
        NodeInfoDef WillAdd = new NodeInfoDef();
        WillAdd.ParentNode = pNode;
        WillAdd.Rank = 0;
        mNodeInfoDict[pNode] = WillAdd;
    }

    // 合併処理
    internal void Unite(int pX, int pY)
    {
        int XNode = FindSet(pX);
        int YNode = FindSet(pY);
        int XRank = mNodeInfoDict[XNode].Rank;
        int YRank = mNodeInfoDict[YNode].Rank;

        if (XRank > YRank) {
            mNodeInfoDict[YNode].ParentNode = XNode;
        }
        else {
            mNodeInfoDict[XNode].ParentNode = YNode;
            if (XRank == YRank) {
                mNodeInfoDict[YNode].Rank++;
            }
        }
    }

    // ノードを引数として、木の根を取得
    internal int FindSet(int pTargetNode)
    {
        // 根までの経路上のノードのList
        var PathNodeList = new List<int>();

        int CurrNode = pTargetNode;
        while (CurrNode != mNodeInfoDict[CurrNode].ParentNode) {
            PathNodeList.Add(CurrNode);
            CurrNode = mNodeInfoDict[CurrNode].ParentNode;
        }

        // 経路圧縮 (親ポインタの付け替え)
        foreach (int EachPathNode in PathNodeList) {
            mNodeInfoDict[EachPathNode].ParentNode = CurrNode;
        }
        return CurrNode;
    }

    internal void DebugPrint()
    {
        foreach (var EachPair in mNodeInfoDict.OrderBy(pX => pX.Key)) {
            Console.WriteLine("mNodeInfoDict[{0}].ParentNode={1}",
                EachPair.Key, EachPair.Value.ParentNode);
        }
    }
}
#endregion

#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)
    {
        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)
    {
        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)
    {
        int UpperB = UpperBound(v);
        if (IsValidInd(UpperB - 1)) {
            return UpperB - 1;
        }
        return -1;
    }

    // 追加機能 V以下で最大の要素の番号を返す
    internal int LowerOrEqual_Max(T v)
    {
        int LowerB = LowerBound(v);
        if (IsValidInd(LowerB - 1)) {
            return LowerB - 1;
        }
        return -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

解説

UnionFindで森を管理しつつ、
木ごとの残次数と
木ごとの接続可能ノード
などを管理してます。

残次数が最小の木と
残次数が最大の木の
接続を繰り返して解いてます。