ABC306

Published on
Updated on

はじめに

AtCoder Beginner Contest 306の復習記事です。

記事におけるScannerクラスは、自作の入力クラスです。

Scannerクラス
public static class Scanner
{
    public static T Scan<T>() where T : IConvertible => Convert<T>(ScanStringArray()[0]);
    public static (T1, T2) Scan<T1, T2>() where T1 : IConvertible where T2 : IConvertible
    {
        var input = ScanStringArray();
        return (Convert<T1>(input[0]), Convert<T2>(input[1]));
    }
    public static (T1, T2, T3) Scan<T1, T2, T3>() where T1 : IConvertible where T2 : IConvertible where T3 : IConvertible
    {
        var input = ScanStringArray();
        return (Convert<T1>(input[0]), Convert<T2>(input[1]), Convert<T3>(input[2]));
    }
    public static (T1, T2, T3, T4) Scan<T1, T2, T3, T4>() where T1 : IConvertible where T2 : IConvertible where T3 : IConvertible where T4 : IConvertible
    {
        var input = ScanStringArray();
        return (Convert<T1>(input[0]), Convert<T2>(input[1]), Convert<T3>(input[2]), Convert<T4>(input[3]));
    }
    public static (T1, T2, T3, T4, T5) Scan<T1, T2, T3, T4, T5>() where T1 : IConvertible where T2 : IConvertible where T3 : IConvertible where T4 : IConvertible where T5 : IConvertible
    {
        var input = ScanStringArray();
        return (Convert<T1>(input[0]), Convert<T2>(input[1]), Convert<T3>(input[2]), Convert<T4>(input[3]), Convert<T5>(input[4]));
    }
    public static (T1, T2, T3, T4, T5, T6) Scan<T1, T2, T3, T4, T5, T6>() where T1 : IConvertible where T2 : IConvertible where T3 : IConvertible where T4 : IConvertible where T5 : IConvertible where T6 : IConvertible
    {
        var input = ScanStringArray();
        return (Convert<T1>(input[0]), Convert<T2>(input[1]), Convert<T3>(input[2]), Convert<T4>(input[3]), Convert<T5>(input[4]), Convert<T6>(input[5]));
    }
    public static IEnumerable<T> ScanEnumerable<T>() where T : IConvertible => ScanStringArray().Select(Convert<T>);
    private static string[] ScanStringArray()
    {
        var line = Console.ReadLine()?.Trim() ?? string.Empty;
        return string.IsNullOrEmpty(line) ? Array.Empty<string>() : line.Split(' ');
    }
    private static T Convert<T>(string value) where T : IConvertible => (T)System.Convert.ChangeType(value, typeof(T));
}

コンテスト

https://atcoder.jp/contests/abc306

問題A

コンテスト提出

各文字を2回繰り返したものを出力します。

public static void Solve()
{
    var N = Scanner.Scan<int>();
    var S = Scanner.Scan<string>();
    var builder = new StringBuilder();
    foreach (var c in S)
    {
        builder.Append($"{c}{c}");
    }

    var answer = builder.ToString();
    Console.WriteLine(answer);
}

問題B

コンテスト提出

1iビットシフトすると、2^iになるので、A[i]<<iを足していくことで答えを求めることができます。
符号あり64bit整数型のlong型だと、2^63-1までしか格納できないので、ulong型などの符号なし64bit整数型を使うことに注意が必要です(1敗)。

public static void Solve()
{
    var A = Scanner.ScanEnumerable<int>().ToArray();
    ulong answer = 0;
    for (var i = 0; i < A.Length; i++)
    {
        answer |= (ulong)A[i] << i;
    }

    Console.WriteLine(answer);
}

問題C

コンテスト提出

2回目に出現する値からなる数列が答えとなります。
各値が何回出現したかを管理しながらAを順に走査し、値が2回目に出現したものを答えの数列に追加することで答えを求めることができます。

public static void Solve()
{
    var N = Scanner.Scan<int>();
    var A = Scanner.ScanEnumerable<int>().ToArray();
    var count = new int[N + 1];
    var answer = new List<int>(N);
    foreach (var a in A)
    {
        count[a]++;
        if (count[a] == 2) answer.Add(a);
    }

    Console.WriteLine(string.Join(" ", answer));
}

問題D

コンテスト提出

次のような動的計画法を解きます。

dp[i][f] := i番目の料理まで見たとき、高橋くんがおなかの状態(f=壊している|壊していない)のときの食べた料理のおいしさの総和の最大値

遷移としては次のようになります。

f=0 // おなかを壊していない
f=1 // おなかを壊している

i番目の料理を下げてもらうとき
dp[i+1][0] = Max(dp[i+1][0], dp[i][0])
dp[i+1][1] = Max(dp[i+1][1], dp[i][1])

i番目の解毒剤入りの料理を食べるとき(x==0)
dp[i+1][0] = Max(dp[i+1][0], dp[i][0]+y, dp[i][1]+y)

i番目の毒入りの料理を食べるとき(y==0)
dp[i+1][1] = Max(dp[i+1][1], dp[i][0]+y)

N番目の料理まで見たときのおなかを壊していないときとおなかを壊しているときのうち、最大値が答えとなります。

public static void Solve()
{
    var N = Scanner.Scan<int>();
    var dp = new long[N + 1, 2];
    const int Inf = (int)1e9;
    dp[0, 1] = -Inf;

    for (var i = 0; i < N; i++)
    {
        var (x, y) = Scanner.Scan<int, int>();
        dp[i + 1, 0] = Math.Max(dp[i + 1, 0], dp[i, 0]);
        dp[i + 1, 1] = Math.Max(dp[i + 1, 1], dp[i, 1]);

        if (x == 0)
        {
            dp[i + 1, 0] = Math.Max(dp[i + 1, 0], dp[i, 0] + y);
            dp[i + 1, 0] = Math.Max(dp[i + 1, 0], dp[i, 1] + y);
        }
        else
        {
            dp[i + 1, 1] = Math.Max(dp[i + 1, 1], dp[i, 0] + y);
        }
    }

    var answer = Math.Max(dp[N, 0], dp[N, 1]);
    Console.WriteLine(answer);
}

問題E

コンテスト提出

クエリごとにAに対してソートを行い、降順K個の和を計算してしまうと、クエリ当たりの時間計算量がO(NlogN)、全体の時間計算量がO(QNlogN)となり、実行時間制限に間に合いません。

f(A)Aのうちの降順K個の総和であり、K個の中からある値vが引かれたとき、f(A)からvが引かれ、K+1番目の値がf(A)に加えられます。 対して、K個の中にある値vが加えられたとき、f(A)にはvが加えられ、K+1番目の値がf(A)から引かれます。

つまり、更新前のA[x]K番目以内であれば、f(A)からA[x]を引き、AK+1番目の値をf(A)に加えます。
A[x]=yとしてAを更新後、A[x]K番以内であれば、f(A)A[x]を加え、AK+1番目の値をf(A)から引きます。

現在のAを管理しながらf(A)を更新していくことで答えを求めることができます。 クエリにおいて、現在のAの中からK番目の値を高速に求める必要があり、重複を許す順序集合を使うことで、時間計算量O(logN)で値を取得することができます。

public static void Solve()
{
    var (N, K, Q) = Scanner.Scan<int, int, int>();
    var A = new long[N + 1];
    var set = new Set<long>((x, y) => y.CompareTo(x), true);
    for (var i = 0; i <= N; i++)
    {
        set.Add(0);
    }

    long answer = 0;
    for (var i = 0; i < Q; i++)
    {
        var (x, y) = Scanner.Scan<int, long>();
        var lb1 = set.LowerBound(A[x]);
        if (lb1 < K)
        {
            answer -= A[x];
            answer += set.ElementAt(K);
        }

        set.Remove(A[x]);

        A[x] = y;
        set.Add(y);

        var lb2 = set.LowerBound(y);
        if (lb2 < K)
        {
            answer += y;
            answer -= set.ElementAt(K);
        }

        Console.WriteLine(answer);
    }
}
Set
public class Set<T> : IReadOnlyCollection<T>
{
    private readonly RandomizedBinarySearchTree<T> _tree;
    private readonly bool _allowDuplication;
    public Set(bool allowDuplication = false) : this(Comparer<T>.Default, allowDuplication) { }

    public Set(IEnumerable<T> source, bool allowDuplication = false) : this(allowDuplication)
    {
        foreach (var value in source) Add(value);
    }

    public Set(IEnumerable<T> source, IComparer<T> comparer, bool allowDuplication = false)
        : this(comparer, allowDuplication)
    {
        foreach (var value in source) Add(value);
    }

    public Set(IEnumerable<T> source, Comparison<T> comparison, bool allowDuplication = false)
        : this(comparison, allowDuplication)
    {
        foreach (var value in source) Add(value);
    }

    public Set(IComparer<T> comparer, bool allowDuplication = false)
        : this((comparer ?? Comparer<T>.Default).Compare, allowDuplication)
    {
    }

    public Set(Comparison<T> comparison, bool allowDuplication = false)
    {
        _tree = new RandomizedBinarySearchTree<T>(comparison);
        _allowDuplication = allowDuplication;
    }

    public void Add(T value)
    {
        if (_allowDuplication || !_tree.Contains(value)) _tree.Insert(value);
    }

    public void Remove(T value)
    {
        _tree.Remove(value);
    }

    public bool Contains(T value)
    {
        return _tree.Contains(value);
    }

    public T ElementAt(int index)
    {
        return _tree.ElementAt(index);
    }

    public int LowerBound(T value)
    {
        return _tree.LowerBound(value);
    }

    public int UpperBound(T value)
    {
        return _tree.UpperBound(value);
    }

    public int Count => _tree.Count;
    public IEnumerator<T> GetEnumerator() => _tree.GetEnumerator();
    IEnumerator IEnumerable.GetEnumerator() => GetEnumerator();
}
Randomized Binary Search Tree

public class RandomizedBinarySearchTree<T> : IReadOnlyCollection<T>
{
    private readonly Comparison<T> _comparison;
    private readonly Random _random;
    private Node _root;
    public RandomizedBinarySearchTree(int seed = 0) : this(comparer: null, seed) { }

    public RandomizedBinarySearchTree(Comparer<T> comparer, int seed = 0) : this(
        (comparer ?? Comparer<T>.Default).Compare, seed)
    {
    }

    public RandomizedBinarySearchTree(Comparison<T> comparison, int seed = 0)
    {
        _comparison = comparison;
        _random = new Random(seed);
    }

    public void Insert(T value)
    {
        if (_root is null) _root = new Node(value);
        else InsertAt(LowerBound(value), value);
    }

    public bool Remove(T value)
    {
        var index = LowerBound(value);
        if (index < 0) return false;
        RemoveAt(index);
        return true;
    }

    public T ElementAt(int index)
    {
        if (index < 0 || Count <= index) throw new ArgumentOutOfRangeException(nameof(index));
        var node = _root;
        var idx = CountOf(node) - CountOf(node.R) - 1;
        while (node is { })
        {
            if (idx == index) return node.Value;
            if (idx > index)
            {
                node = node.L;
                idx -= CountOf(node?.R) + 1;
            }
            else
            {
                node = node.R;
                idx += CountOf(node?.L) + 1;
            }
        }

        throw new ArgumentOutOfRangeException(nameof(index));
    }

    public bool Contains(T value)
    {
        return Find(value) is { };
    }

    public int Count => CountOf(_root);
    public IEnumerator<T> GetEnumerator() => Enumerate(_root).GetEnumerator();
    IEnumerator IEnumerable.GetEnumerator() => GetEnumerator();
    public int UpperBound(T value) => Bound(value, (x, y) => _comparison(x, y) > 0);
    public int LowerBound(T value) => Bound(value, (x, y) => _comparison(x, y) >= 0);

    public int Bound(T value, Func<T, T, bool> compare)
    {
        var node = _root;
        if (node is null) return -1;
        var bound = CountOf(node);
        var idx = bound - CountOf(node.R) - 1;
        while (node is { })
        {
            if (compare(node.Value, value))
            {
                node = node.L;
                bound = Math.Min(bound, idx);
                idx -= CountOf(node?.R) + 1;
            }
            else
            {
                node = node.R;
                idx += CountOf(node?.L) + 1;
            }
        }

        return bound;
    }

    private double GetProbability() => _random.NextDouble();

    private void InsertAt(int index, T value)
    {
        var (l, r) = Split(_root, index);
        _root = Merge(Merge(l, new Node(value)), r);
    }

    private void RemoveAt(int index)
    {
        var (l, r1) = Split(_root, index);
        var (_, r2) = Split(r1, 1);
        _root = Merge(l, r2);
    }

    private Node Merge(Node l, Node r)
    {
        if (l is null || r is null) return l ?? r;
        var (n, m) = (CountOf(l), CountOf(r));
        if ((double)n / (n + m) > GetProbability())
        {
            l.R = Merge(l.R, r);
            return l;
        }
        else
        {
            r.L = Merge(l, r.L);
            return r;
        }
    }

    private (Node, Node) Split(Node node, int k)
    {
        if (node is null) return (null, null);
        if (k <= CountOf(node.L))
        {
            var (l, r) = Split(node.L, k);
            node.L = r;
            return (l, node);
        }
        else
        {
            var (l, r) = Split(node.R, k - CountOf(node.L) - 1);
            node.R = l;
            return (node, r);
        }
    }

    private Node Find(T value)
    {
        var node = _root;
        while (node is { })
        {
            var cmp = _comparison(node.Value, value);
            if (cmp > 0) node = node.L;
            else if (cmp < 0) node = node.R;
            else break;
        }

        return node;
    }

    private static int CountOf(Node node) => node?.Count ?? 0;

    private static IEnumerable<T> Enumerate(Node node = null)
    {
        if (node is null) yield break;
        foreach (var value in Enumerate(node.L)) yield return value;
        yield return node.Value;
        foreach (var value in Enumerate(node.R)) yield return value;
    }

    private class Node
    {
        internal T Value { get; }

        internal Node L
        {
            get => _l;
            set
            {
                _l = value;
                UpdateCount();
            }
        }

        internal Node R
        {
            get => _r;
            set
            {
                _r = value;
                UpdateCount();
            }
        }

        internal int Count { get; private set; }
        private Node _l;
        private Node _r;

        internal Node(T value)
        {
            Value = value;
            Count = 1;
        }

        private void UpdateCount()
        {
            Count = (L?.Count ?? 0) + (R?.Count ?? 0) + 1;
        }
    }
}