Being Code Expert: The chocolate rooms

Your task is to take some chocolates from N different rooms. Each room contains P[i] number of chocolates of the same or different brands. You are given an integer K.

Write a program to determine whether you can visit each room and find K different brands of chocolates.

Example:

Input:  n=3, k=2, arr=[1,2,1], s={{"KITKAT"},{"FIVESTAR","KITKAT"},{"KITKAT"}}
Output: Yes

Approach

Java

import java.util.HashSet;

public class ChocolateRooms {
    public static void main(String[] args) {

        int n = 3, k = 2;

        int[] arr = { 1, 2, 1 };

        String[][] s = { { "KITKAT" }, { "FIVESTAR", "KITKAT" },
        { "KITKAT" } };

        System.out.println(chocolateRooms(arr, s, n, k));

    }

    static String chocolateRooms(int[] arr, String[][] s,
       int n, int k)

    {
        HashSet<String> st = new HashSet<String>();
        for (int i = 0; i < n; i++) {
            int p = arr[i];

            for (int j = 0; j < p; j++) {

                st.add(s[i][j]);
            }
        }
        if (st.size() >= k)
            return "Yes";
        else
            return "No";
    }

}

C++

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

string chocolateRooms(vector<int> &arr, 
      vector<vector<string>> &s,
                      int n, int k)

{
    set<string> st;
    for (int i = 0; i < n; i++)
    {
        int p = arr[i];

        for (int j = 0; j < p; j++)
        {

            st.insert(s[i][j]);
        }
    }
    if (st.size() >= k)
        return "Yes";
    else
        return "No";
}
int main()
{
    int n = 3, k = 2;

    vector<int> arr = {1, 2, 1};

    vector<vector<string>> s = {{"KITKAT"},
                                {"FIVESTAR", "KITKAT"},
                                {"KITKAT"}};

    cout << chocolateRooms(arr, s, n, k) << "\n";

    return 0;
}