mirror of
https://github.com/tgorordo/carousel.git
synced 2026-06-05 18:12:14 -07:00
80 lines
1.9 KiB
Python
80 lines
1.9 KiB
Python
import polars as pl
|
|
from polars.testing import assert_frame_equal
|
|
from carousel import GS_deferred_acceptance
|
|
|
|
|
|
def test_prefs():
|
|
people_prefs = pl.DataFrame(
|
|
{
|
|
"people": [
|
|
"Alice",
|
|
"Alice",
|
|
"Alice",
|
|
"Bob",
|
|
"Bob",
|
|
"Bob",
|
|
"Charlie",
|
|
"Charlie",
|
|
"Charlie",
|
|
],
|
|
"fruit": [
|
|
"apple",
|
|
"banana",
|
|
"cherry",
|
|
"banana",
|
|
"cherry",
|
|
"apple",
|
|
"cherry",
|
|
"apple",
|
|
"banana",
|
|
],
|
|
"rank": [1, 2, 3, 1, 2, 3, 1, 2, 3],
|
|
}
|
|
)
|
|
|
|
fruit_prefs = pl.DataFrame(
|
|
{
|
|
"fruit": [
|
|
"apple",
|
|
"apple",
|
|
"apple",
|
|
"banana",
|
|
"banana",
|
|
"banana",
|
|
"cherry",
|
|
"cherry",
|
|
"cherry",
|
|
],
|
|
"people": [
|
|
"Alice",
|
|
"Bob",
|
|
"Charlie",
|
|
"Alice",
|
|
"Bob",
|
|
"Charlie",
|
|
"Alice",
|
|
"Bob",
|
|
"Charlie",
|
|
],
|
|
"rank": [1, 1, 1, 1, 1, 1, 1, 1, 1], # fruits have no preferences
|
|
}
|
|
)
|
|
|
|
capacities = pl.DataFrame(
|
|
{
|
|
"fruit": ["apple", "cherry", "banana"],
|
|
"capacity": [1, 1, 1], # have one of each
|
|
}
|
|
)
|
|
|
|
assert_frame_equal(
|
|
GS_deferred_acceptance(
|
|
people_prefs, fruit_prefs, capacities, app_col="people", pos_col="fruit"
|
|
).sort(["people", "fruit"]),
|
|
pl.DataFrame(
|
|
{
|
|
"people": ["Alice", "Bob", "Charlie"],
|
|
"fruit": ["apple", "banana", "cherry"],
|
|
}
|
|
).sort(["people", "fruit"]),
|
|
)
|