Implement Table Aggregation
Company: Notion
Role: Software Engineer
Category: Coding & Algorithms
Difficulty: easy
Interview Round: Technical Screen
You are given a table represented as a list of rows. Each row contains:
- `key`: a string identifying a group
- `value`: an integer
Write a function that aggregates the table by `key` and returns, for each distinct key, the sum of all corresponding `value`s.
Example:
Input rows:
- `(A, 3)`
- `(B, 5)`
- `(A, 2)`
- `(B, 1)`
- `(C, 4)`
Output:
- `(A, 5)`
- `(B, 6)`
- `(C, 4)`
The output may be returned in any order unless otherwise specified.
Follow-up discussion:
1. How would you extend the implementation so that it can support other aggregation functions such as `MAX` in addition to `SUM`?
2. If the table is too large to fit into memory, how would you process it efficiently?
Overview: This question evaluates a candidate's ability to perform data aggregation and manipulate basic data structures by grouping records by key and combining numeric values.
You are given a table represented as a list of rows. Each row is a pair `[key, value]` where `key` is a string identifying a group and `value` is an integer. Aggregate the table by `key` and return, for each distinct key, the sum of all corresponding values.
Return the result as a list of `[key, sum]` pairs sorted in ascending order by key (the underlying problem allows any order, but this console fixes a sorted order so the result is deterministic).
Example:
Input rows: `[["A", 3], ["B", 5], ["A", 2], ["B", 1], ["C", 4]]`
Output: `[["A", 5], ["B", 6], ["C", 4]]`
Follow-up discussion (not graded here): How would you generalize the implementation to support other aggregations such as `MAX` in addition to `SUM`? And how would you process the table efficiently if it is too large to fit in memory (e.g. external/streaming aggregation, or map-reduce style partial sums by key)?
Constraints
- 0 <= number of rows <= 10^5
- Each row is a [key, value] pair.
- key is a non-empty string.
- value is an integer that may be negative, zero, or positive.
- Sums fit in a 64-bit integer.
Examples
Input: [["A", 3], ["B", 5], ["A", 2], ["B", 1], ["C", 4]]
Expected Output: [["A", 5], ["B", 6], ["C", 4]]
Explanation: A: 3+2=5, B: 5+1=6, C: 4. Sorted by key.
Input: []
Expected Output: []
Explanation: Empty table aggregates to an empty result.
Hints
- Use a hash map keyed by the group string, adding each row's value into the running total for that key.
- Collect the distinct keys and sort them so the output order is deterministic.
- For the MAX follow-up, replace the '+=' accumulation with a reducer function (e.g. max) so the same loop supports any associative aggregation. For the out-of-memory follow-up, stream the rows and aggregate incrementally, or partition by key and aggregate each partition (map-reduce).
Community answers
Answer by thedevunit
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
public class TableAggregation {
public record Row(String key, int value) {}
public static List aggregateBySum(List rows) {
Map sums = new HashMap<>();
for (Row row : rows) {
sums.merge(row.key(), row.value(), Integer::sum);
}
List result = new ArrayList<>();
for (Map.Entry entry : sums.entrySet()) {
result.add(new Row(entry.getKey(), entry.getValue()));
}
return result;
}
public static void main(String[] args) {
List rows = List.of(
new Row("A", 3),
new Row("B", 5),
new Row("A", 2),
new Row("B", 1),
new Row("C", 4)
);
System.out.println(aggregateBySum(rows));
}
}
Pass the aggregation logic as a function instead of hardcoding SUM.
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.function.IntBinaryOperator;
public class TableAggregator {
public record Row(String key, int value) {}
public static Map aggregate(
List rows,
IntBinaryOperator aggregationFunction) {
Map result = new HashMap<>();
for (Row row : rows) {
result.merge(
row.key(),
row.value(),
aggregationFunction::applyAsInt
);
}
return result;
}
}
Now the caller can choose the aggregation:
Map sums =
aggregate(rows, Integer::sum);
Map maximums =
aggregate(rows, Math::max);
Map minimums =
aggregate(rows, Math::min);