Find the Number of Distinct Colors Among the Balls
Problem
There are limit + 1 balls labeled 0..limit, all initially uncolored. Each query [x, y] repaints ball x with color y (overwriting any previous color). After each query, report how many distinct colors are currently in use. Return the array of answers, one per query. Uncolored balls do not count as a color.
limit = 4, queries = [[1,4],[2,5],[1,3],[3,4]][1, 2, 2, 3]def queryResults(limit, queries):
ballColor = {} # ball -> its current color
colorCount = {} # color -> how many balls wear it
result = []
for x, y in queries:
if x in ballColor: # ball was colored before
old = ballColor[x]
colorCount[old] -= 1 # that color loses one ball
if colorCount[old] == 0:
del colorCount[old] # color no longer present
ballColor[x] = y # repaint ball x with color y
colorCount[y] = colorCount.get(y, 0) + 1
result.append(len(colorCount)) # distinct colors right now
return resultfunction queryResults(limit, queries) {
const ballColor = new Map(); // ball -> its current color
const colorCount = new Map(); // color -> how many balls wear it
const result = [];
for (const [x, y] of queries) {
if (ballColor.has(x)) { // ball was colored before
const old = ballColor.get(x);
colorCount.set(old, colorCount.get(old) - 1);
if (colorCount.get(old) === 0) // color no longer present
colorCount.delete(old);
}
ballColor.set(x, y); // repaint ball x with color y
colorCount.set(y, (colorCount.get(y) || 0) + 1);
result.push(colorCount.size); // distinct colors right now
}
return result;
}int[] queryResults(int limit, int[][] queries) {
Map<Integer, Integer> ballColor = new HashMap<>(); // ball -> color
Map<Integer, Integer> colorCount = new HashMap<>(); // color -> count
int[] result = new int[queries.length];
for (int i = 0; i < queries.length; i++) {
int x = queries[i][0], y = queries[i][1];
if (ballColor.containsKey(x)) { // colored before
int old = ballColor.get(x);
colorCount.merge(old, -1, Integer::sum); // that color -1
if (colorCount.get(old) == 0)
colorCount.remove(old); // color gone
}
ballColor.put(x, y); // repaint ball x
colorCount.merge(y, 1, Integer::sum);
result[i] = colorCount.size(); // distinct colors
}
return result;
}vector<int> queryResults(int limit, vector<vector<int>>& queries) {
unordered_map<int, int> ballColor; // ball -> its color
unordered_map<int, int> colorCount; // color -> count
vector<int> result;
for (auto& q : queries) {
int x = q[0], y = q[1];
auto it = ballColor.find(x);
if (it != ballColor.end()) { // colored before
int old = it->second;
if (--colorCount[old] == 0) // that color loses a ball
colorCount.erase(old); // color no longer present
}
ballColor[x] = y; // repaint ball x
colorCount[y]++;
result.push_back(colorCount.size()); // distinct colors now
}
return result;
}Explanation
The naive answer would re-scan every ball after each query, but that is too slow. Instead we keep two hash maps and update them in O(1) per query.
ballColor maps each ball to its current color. colorCount maps each color to the number of balls currently wearing it. The answer for any query is simply the number of keys in colorCount — that is how many distinct colors are in use.
When a query [x, y] arrives we first check whether ball x already had a color. If it did, that old color loses one wearer, so we decrement its count; if the count drops to zero we delete the key, because a color with no balls is not a color anymore.
Then we paint ball x with y: record it in ballColor, and increment colorCount[y]. After the update, len(colorCount) is the distinct-color count we append to the result.
Example: queries = [[1,4],[2,5],[1,3],[3,4]]. Query 3 repaints ball 1 from color 4 to color 3. Color 4 had only ball 1, so it disappears as color 3 appears — the distinct count stays at 2, giving [1, 2, 2, 3].