You want to group elements that have at least one common group value ("country" or "age_group", but this applies to any number of grouping values). This can be done for instance with the help of a class that makes building such groups:
class MultiLabelGroup:
"""
Group of indexes pointing to data that belongs to the same group.
Two data samples belong to the same group if they share any value
among multiple labeled values (typically in some columns of a
Pandas DataFrame).
Notable attributes:
- indexes: set with the indexes of the element in the group.
"""
def __init__(self, monitored_indexes):
"""
monitored_indexes -- sequence with the Pandas Series index to monitor
for building the group.
"""
# Sample indexes in the group:
self.indexes = set()
# Values of the columns in the samples of the group:
self._values = {monitored_index: set() for monitored_index in monitored_indexes}
def __repr__(self):
return repr(self.indexes)
def belongs_to_group(self, series: pd.Series) -> bool:
"""
Return whether the sample in the given series belongs to the group.
series -- series with the monitored names (indexes) in its index. The values
of the series at these positions determine whether the sample
belongs to the group.
"""
return any(series[index] in self._values[index] for index in self._values)
def add_to_group(self, index, series: pd.Series):
"""
Add the given index and sample to the group.
index -- index value (user-defined) of the given series (not of a value in the series!).
It must be hashable.
"""
self.indexes.add(index)
for (name, values) in self._values.items():
values.add(series[name])
def merge(self, other_group):
"""
Merge the current group with another group.
other_group -- another MultiLabelGroup, with the same monitored
indexes.
"""
self.indexes |= other_group.indexes
for name in self._values:
self._values[name] |= other_group._values[name]
This group class can be used like so:
groups = []
for index, row in df.iterrows():
groups_to_merge = []
other_groups = []
for group in groups:
(groups_to_merge if group.belongs_to_group(row) else other_groups).append(group)
if groups_to_merge:
# All the matching groups are linked
# by the new data, which must be in all these
# groups. Therefore they must be merged into
# a single new group:
new_group = groups_to_merge.pop()
for group in groups_to_merge:
new_group.merge(group)
groups = other_groups
else: # No existing group
new_group = MultiLabelGroup(["country", "age_group"])
groups.append(new_group)
# The new data must be registered:
groups[-1].add_to_group(index, row)
print(groups)
You get the following single group:
[{0, 1, 2, 3, 4, 5, 6, 7}]
where you see that indeed the element you quoted (element 6 with country "c" and age group "e") is in the same (evaluation/test) set as element 0 (because it has the same age group "e").
The result shows that the data in the original question cannot be split between training and evaluation ("test") sets that do not share either a country or an age group.
This may work for some other, larger datasets, though.
