PG Phriday: Instant Runoff Through SQL

CREATE TABLE irv_election
(
  election_id    SERIAL PRIMARY KEY,
  election_name  VARCHAR NOT NULL
);

CREATE INDEX idx_election_election_name
    ON irv_election (election_name);

CREATE TABLE irv_candidate
(
  candidate_id    SERIAL PRIMARY KEY,
  election_id     INT NOT NULL REFERENCES irv_election,
  candidate_name  VARCHAR NOT NULL
);

CREATE INDEX idx_candidate_candidate_name
    ON irv_candidate (candidate_name);

CREATE TABLE irv_vote
(
  vote_id         SERIAL PRIMARY KEY,
  election_id     INT NOT NULL REFERENCES irv_election,
  candidate_list  INT[] NOT NULL
);

CREATE INDEX idx_vote_election_id
    ON irv_vote (election_id);

CREATE TABLE irv_scratch
(
  candidate_name  VARCHAR,
  votes           BIGINT,
  percent         NUMERIC,
  my_round        INT
);

CREATE OR REPLACE FUNCTION start_election(
  election VARCHAR
)
RETURNS VOID AS
$$
  INSERT INTO irv_election (election_name)
  VALUES (election);
$$ LANGUAGE SQL;


CREATE OR REPLACE FUNCTION register_candidate(
  election VARCHAR,
  candidate VARCHAR
)
RETURNS VOID AS
$$
  INSERT INTO irv_candidate (election_id, candidate_name)
  SELECT election_id, candidate
    FROM irv_election
   WHERE election_name = election;
$$ LANGUAGE SQL;


CREATE OR REPLACE FUNCTION register_vote(
  election VARCHAR,
  candidates VARCHAR[]
)
RETURNS VOID AS
$$
  WITH ranked AS (
    SELECT candidate_name, row_number() OVER () AS rank
      FROM unnest(candidates) candidate_name
  ),
  conv AS (
    SELECT c.election_id, c.candidate_id, r.rank
      FROM irv_candidate c
      JOIN ranked r USING (candidate_name)
     ORDER BY rank
  )
  INSERT INTO irv_vote (election_id, candidate_list)
  SELECT election_id, array_agg(candidate_id)
    FROM conv
   GROUP BY 1;
$$ LANGUAGE SQL;

CREATE OR REPLACE FUNCTION no_losers(orig INT[], remove INT[])
RETURNS INT[] AS
$$
DECLARE
  item INT; 
  ret_arr INT[] := orig;
BEGIN
  IF array_length(remove, 1) IS NULL THEN
    RETURN ret_arr;
  END IF;

  FOR item IN 1 .. array_length(remove, 1) LOOP
    ret_arr := ret_arr - remove[item];
  END LOOP;

  RETURN ret_arr;
END;
$$ LANGUAGE PLPGSQL IMMUTABLE STRICT;


CREATE OR REPLACE FUNCTION tabulate_votes(
  election VARCHAR
)
RETURNS VOID AS
$$
DECLARE
  num_candidates INT;
  num_votes NUMERIC;
  vote_round INT;
  losers INT[] := '{}'::INT[]; -- Non-null empty array.
  winner VARCHAR;
  results RECORD;
BEGIN

  DROP TABLE IF EXISTS working_round;
  CREATE TABLE working_round (LIKE irv_scratch);

  -- Given the number of candidates in this election, we can
  -- reduce calculations.

  SELECT count(*) INTO num_candidates
    FROM irv_candidate c
    JOIN irv_election e USING (election_id)
   WHERE e.election_name = election;

  -- We need the total number of votes cast for display
  -- percentages.

  SELECT count(*) INTO num_votes
    FROM irv_vote v
    JOIN irv_election e USING (election_id)
   WHERE e.election_name = election;

  -- Record each round of the IRV process. At the end of each round,
  -- eliminate the worst performer and try again. Do this until we
  -- reach > 50% for one candidate.

  FOR vote_round IN 1 .. (num_candidates - 1) LOOP
    RAISE NOTICE '== Round % ==', vote_round;

    INSERT INTO working_round
    SELECT c.candidate_name, count(v.vote_id) AS votes,
           round(count(v.vote_id) / 
                 num_votes * 100.0, 2) AS percent,
           vote_round AS my_round
      FROM irv_vote v
      JOIN irv_candidate c ON (
             c.candidate_id = (no_losers(v.candidate_list, losers))[1]
           )
      JOIN irv_election e ON (e.election_id = v.election_id)
     WHERE e.election_name = election
     GROUP BY c.candidate_name;

    -- Print the results of the round for spot-checking.

    FOR results IN
        SELECT * FROM working_round
         WHERE my_round = vote_round
    LOOP
      RAISE NOTICE '%: % (%)', results.candidate_name,
        results.votes, results.percent;
    END LOOP;

    -- If this round has a winner, short-circuit so we can
    -- just report the results.

    SELECT candidate_name INTO winner
      FROM working_round
     WHERE percent > 50
       AND my_round = vote_round;

    IF winner IS NOT NULL THEN
      RAISE NOTICE 'Winner of % is %!', election, winner;
      EXIT;
    END IF;

    -- This is where we determine the loser of this round.
    -- It's just the lowest ranked result.

    SELECT array_append(losers, c.candidate_id) INTO losers
      FROM working_round w
      JOIN irv_candidate c USING (candidate_name)
     WHERE my_round = vote_round
     ORDER BY w.votes
     LIMIT 1;

  END LOOP;

END;
$$ LANGUAGE PLPGSQL;

SELECT start_election('Pres 2016');

SELECT register_candidate('Pres 2016', 'Clinton');
SELECT register_candidate('Pres 2016', 'Trump');
SELECT register_candidate('Pres 2016', 'Stein');
SELECT register_candidate('Pres 2016', 'Johnson');

SELECT register_vote('Pres 2016', '{Clinton, Stein}')
  FROM generate_series(1, 60467);

SELECT register_vote('Pres 2016', '{Stein, Clinton}')
  FROM generate_series(1, 1237);

SELECT register_vote('Pres 2016', '{Trump, Johnson}')
  FROM generate_series(1, 60071);

SELECT register_vote('Pres 2016', '{Johnson, Trump}')
  FROM generate_series(1, 4123);

SELECT tabulate_votes('Pres 2016');

NOTICE:  == Round 1 ==
NOTICE:  Clinton: 60467 (48.03)
NOTICE:  Stein: 1237 (0.98)
NOTICE:  Trump: 60071 (47.71)
NOTICE:  Johnson: 4123 (3.27)
NOTICE:  == Round 2 ==
NOTICE:  Clinton: 61704 (49.01)
NOTICE:  Trump: 60071 (47.71)
NOTICE:  Johnson: 4123 (3.27)
NOTICE:  == Round 3 ==
NOTICE:  Clinton: 61704 (49.01)
NOTICE:  Trump: 64194 (50.99)
NOTICE:  Winner of Pres 2016 is Trump!