【AtCoder】ABC 402 C - Dislike Foods | 緑コーダーが解くAtCoder

配点: 300 点 / 実行時間制限: 2 sec / メモリ制限: 1024 MB / Difficulty: ???

問題概要
制約
考察
実装例
コメント

問題概要

あるレストランでは $1$ から $N$ までの番号がつけられている $N$ 種類の食材を扱っており、また、 $1$ から $M$ までの番号がつけられている $M$ 個の料理を提供している。料理 $i$ には食材 $A_{i,1}, A_{i,2}, \cdots, A_{i,K_i}$ の $K_i$ 種類の食材が使われている。

すぬけ君は現在 $N$ 種類の食材全てが苦手であり、苦手な食材が1種類でも使われている料理を食べることができない。しかし、すぬけ君はこれから $N$ 日間かけて苦手な食材を克服しようとしており、具体的には $i$ 日目に食材 $B_i$ を克服し、それ以降苦手な食材でなくなる。

このとき、 $i=1,2,\cdots,N$ について以下の値を求めよ。

$i$ 日目にすぬけ君が食材 $B_i$ を克服した直後、すぬけ君が食べることができる料理の個数

制約

$1 \leq N, M \leq 3 \times 10^{5}$
$\sum K_i \leq 3 \times 10^{5}$
入力は全て整数

考察

まず愚直に考えると、

1日ごとに、 $M$ 種類の料理それぞれを見ていき、使われている食材が苦手な食材かをチェックする

という方法が思いつく。しかし、これでは計算量が $O(N(\sum K_i + M))$ となってしまい、 TLE する。

明らかに無駄なのは、毎日全ての料理について食べられるか否かをチェックしていることだ。そこで、 $i$ 日目にチェックする料理は食材 $B_i$ が使われているものだけにしたい。

したがって、以下の2種類の情報を管理することにする。

$\mathrm{foods\_to\_dishes}_{i} :=$ 食材 $i$ が使われている料理のリスト
$\mathrm{foods\_num}_{i} :=$ 料理 $i$ に使われている食材のうち、現時点で苦手なものの個数

その上で、 $i$ 日目には次の操作を行う。

の中の各料理について、
1. $\mathrm{foods\_num}_{\mathrm{dish}_{j}}$ を $1$ 減らす。
2. もし、この日に初めて $\mathrm{foods\_num}_{\mathrm{dish}_{j}} = 0$ となったならば、答えのカウンターを $1$ 増やす。
答えのカウンターを出力する。

この方針であれば、計算量は $O(N + \sum K_i + M)$ となるので十分高速である。

実装例

#include <bits/stdc++.h>
using namespace std;

#define rep(i, start, end) for (auto i = (start); (i) < (end); (i)++)

// ======================================== //

int main()
{
    int N, M;
    cin >> N >> M;
    vector<vector<int>> foods_to_dishes(N + 1);
    vector<int> foods_num(M, 0);
    rep(i, 0, M)
    {
        int K;
        cin >> K;
        rep(j, 0, K)
        {
            int A;
            cin >> A;
            foods_to_dishes[A].push_back(i);
            foods_num[i]++;
        }
    }
    vector<int> B(N);
    rep(i, 0, N) cin >> B[i];

    vector<bool> eatable_dish(M, false);
    int ans = 0;
    rep(i, 0, N)
    {
        for (auto &&dish : foods_to_dishes[B[i]])
        {
            foods_num[dish]--;

            if (foods_num[dish] == 0)
            {
                if (!eatable_dish[dish])
                {
                    ans++;
                    eatable_dish[dish] = true;
                }
            }
        }

        cout << ans << endl;
    }

    return 0;
}