Skip to content

3160. Find the Number of Distinct Colors Among the Balls 👍

  • Time: $O(q)$
  • Space: $O(q)$
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
class Solution {
 public:
  vector<int> queryResults(int limit, vector<vector<int>>& queries) {
    vector<int> ans;
    unordered_map<int, int> ballToColor;
    unordered_map<int, int> colorCount;

    for (const vector<int>& query : queries) {
      const int ball = query[0];
      const int color = query[1];
      if (const auto it = ballToColor.find(ball); it != ballToColor.cend()) {
        const int prevColor = it->second;
        if (--colorCount[prevColor] == 0)
          colorCount.erase(prevColor);
      }
      ballToColor[ball] = color;
      ++colorCount[color];
      ans.push_back(colorCount.size());
    }

    return ans;
  }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
class Solution {
  public int[] queryResults(int limit, int[][] queries) {
    int[] ans = new int[queries.length];
    Map<Integer, Integer> ballToColor = new HashMap<>();
    Map<Integer, Integer> colorCount = new HashMap<>();

    for (int i = 0; i < queries.length; ++i) {
      final int ball = queries[i][0];
      final int color = queries[i][1];
      if (ballToColor.containsKey(ball)) {
        final int prevColor = ballToColor.get(ball);
        if (colorCount.merge(prevColor, -1, Integer::sum) == 0)
          colorCount.remove(prevColor);
      }
      ballToColor.put(ball, color);
      colorCount.merge(color, 1, Integer::sum);
      ans[i] = colorCount.size();
    }

    return ans;
  }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
class Solution:
  def queryResults(self, limit: int, queries: list[list[int]]) -> list[int]:
    ans = []
    ballToColor = {}
    colorCount = collections.Counter()

    for ball, color in queries:
      if ball in ballToColor:
        prevColor = ballToColor[ball]
        colorCount[prevColor] -= 1
        if colorCount[prevColor] == 0:
          del colorCount[prevColor]
      ballToColor[ball] = color
      colorCount[color] += 1
      ans.append(len(colorCount))

    return ans