Add ranked ballot marker filter

src/votekit/cleaning/__init__.py

@@ -8,6 +8,7 @@
     condense_rank_profile,
     remove_and_condense_rank_profile,
     remove_cand_rank_profile,
+    remove_rank_ballots_with_cands,
     remove_repeat_cands_rank_profile,
 )
 from .score_ballots_cleaning import remove_cand_score_ballot
@@ -17,6 +18,7 @@
     "clean_rank_profile",
     "remove_repeat_cands_rank_profile",
     "remove_cand_rank_profile",
+    "remove_rank_ballots_with_cands",
     "condense_rank_profile",
     "remove_and_condense_rank_profile",
     "remove_cand_rank_ballot",

src/votekit/cleaning/rank_profiles_cleaning.py

@@ -288,6 +288,61 @@ def remove_cand_rank_profile(
     )
 
 
+def remove_rank_ballots_with_cands(
+    removed: Candidate | CandidateList,
+    profile: RankProfile,
+    remove_zero_weight_ballots: bool = True,
+) -> CleanedRankProfile:
+    """
+    Remove ballots containing any specified candidate or marker from a ranked profile.
+
+    Args:
+        removed (Candidate | list[Candidate] | list[str] | list[int]):
+            Candidate, marker, or list of candidates and markers to filter on.
+        profile (RankProfile): Profile to remove matching ballots from.
+        remove_zero_weight_ballots (bool, optional): Whether or not to remove ballots that have no
+            weight. Defaults to True.
+
+    Returns:
+        CleanedRankProfile: A cleaned ``RankProfile``.
+
+    Raises:
+        ProfileError: Profile must only contain ranked ballots.
+    """
+    if not isinstance(profile, RankProfile):
+        raise ProfileError("Profile must be a RankProfile.")
+
+    if isinstance(removed, Candidate):
+        removed = [removed]
+
+    removed_set = set(removed)
+
+    assert profile.max_ranking_length is not None
+    ranking_cols = [f"Ranking_{i}" for i in range(1, profile.max_ranking_length + 1)]
+    contains_removed = profile.df[ranking_cols].map(
+        lambda cands: isinstance(cands, frozenset) and not cands.isdisjoint(removed_set)
+    )
+    matching_idxs = set(profile.df.index[contains_removed.any(axis=1)])
+
+    cleaned_df = profile.df.drop(index=matching_idxs)
+    if remove_zero_weight_ballots:
+        cleaned_df = cleaned_df[cleaned_df["Weight"] > 0]
+
+    unaltr_idxs = set(profile.df.index) - matching_idxs
+
+    return CleanedRankProfile(
+        df=cleaned_df,
+        candidates=profile.candidates,
+        max_ranking_length=profile.max_ranking_length,
+        parent_profile=profile,
+        df_index_column=list(cleaned_df.index),
+        no_wt_altr_idxs=set(),
+        no_rank_altr_idxs=matching_idxs,
+        nonempty_altr_idxs=set(),
+        unaltr_idxs=unaltr_idxs,
+    )
+
+
 def condense_ranking_row(
     ranking_tup: tuple,
 ) -> tuple:

tests/cleaning/rank_profiles/test_remove_rank_ballots_with_cands.py

@@ -0,0 +1,65 @@
+import pytest
+
+from votekit.ballot import RankBallot, ScoreBallot
+from votekit.cleaning import remove_rank_ballots_with_cands
+from votekit.pref_profile import CleanedRankProfile, ProfileError, RankProfile, ScoreProfile
+
+profile = RankProfile(
+    ballots=[
+        RankBallot(ranking=[{"A"}, {"B"}], weight=1),
+        RankBallot(ranking=[{"A"}, {"overvote"}, {"C"}], weight=2),
+        RankBallot(ranking=[{"C"}, {"A"}], weight=3),
+        RankBallot(ranking=[{"undervote", "B"}, {"C"}], weight=4),
+        RankBallot(ranking=[{"A"}], weight=0),
+    ]
+)
+
+
+def test_remove_rank_ballots_with_candidate():
+    cleaned_profile = remove_rank_ballots_with_cands("overvote", profile)
+
+    assert isinstance(cleaned_profile, CleanedRankProfile)
+    assert cleaned_profile.parent_profile == profile
+    assert cleaned_profile.ballots == (
+        RankBallot(ranking=[{"A"}, {"B"}], weight=1),
+        RankBallot(ranking=[{"C"}, {"A"}], weight=3),
+        RankBallot(ranking=[{"undervote", "B"}, {"C"}], weight=4),
+    )
+    assert cleaned_profile.candidates == profile.candidates
+    assert cleaned_profile.no_wt_altr_idxs == set()
+    assert cleaned_profile.no_rank_altr_idxs == {1}
+    assert cleaned_profile.nonempty_altr_idxs == set()
+    assert cleaned_profile.unaltr_idxs == {0, 2, 3, 4}
+
+
+def test_remove_rank_ballots_with_multiple_candidates_and_ties():
+    cleaned_profile = remove_rank_ballots_with_cands(["overvote", "undervote"], profile)
+
+    assert cleaned_profile.ballots == (
+        RankBallot(ranking=[{"A"}, {"B"}], weight=1),
+        RankBallot(ranking=[{"C"}, {"A"}], weight=3),
+    )
+    assert cleaned_profile.no_rank_altr_idxs == {1, 3}
+    assert cleaned_profile.unaltr_idxs == {0, 2, 4}
+
+
+def test_remove_rank_ballots_with_candidate_retains_zero_weight_when_requested():
+    cleaned_profile = remove_rank_ballots_with_cands(
+        "overvote",
+        profile,
+        remove_zero_weight_ballots=False,
+    )
+
+    assert cleaned_profile.ballots == (
+        RankBallot(ranking=[{"A"}, {"B"}], weight=1),
+        RankBallot(ranking=[{"C"}, {"A"}], weight=3),
+        RankBallot(ranking=[{"undervote", "B"}, {"C"}], weight=4),
+        RankBallot(ranking=[{"A"}], weight=0),
+    )
+
+
+def test_remove_rank_ballots_with_cands_requires_rank_profile():
+    score_profile = ScoreProfile(ballots=[ScoreBallot(scores={"A": 1})])
+
+    with pytest.raises(ProfileError, match="Profile must be a RankProfile."):
+        remove_rank_ballots_with_cands("A", score_profile)  # type: ignore[arg-type]