用SQL解海盗分金问题

spoils.1.sql

create table spoils (
  id int primary key,
  amount int
);
comment on table spoils is '分赃表';
comment on column spoils.id is '海盗编号';
comment on column spoils.amount is '可获得的最高金额，null代表死亡';

create or replace function divide(amount_total int, pirate_count int)
  returns setof spoils as $$
declare
  _id int;
  _cost int;       -- 用于贿赂其他海盗的成本
  _partners int[]; -- 接受贿赂的海盗同伙
begin
  delete from spoils;
  if pirate_count > 0 then
    for _id in 1..pirate_count loop
      select
        coalesce(sum(coalesce(amount + 1, 0)), 0),
        coalesce(array_agg(id), '{}')
      into
        _cost,
        _partners
      from (
        select
          *
        from
          spoils
        order by
          amount asc nulls first -- 选出贿赂成本最低的
        limit
          _id / 2                -- 只选出半数海盗贿赂即可
      ) as t;

      if _cost <= amount_total then
        update spoils set amount = case
                                     when id = any(_partners)
                                       then coalesce(amount + 1, 0)
                                     else 0
                                   end; -- 没死亡的同伴贿赂金额要上升
        insert into spoils values (_id, amount_total - _cost);
      else
        insert into spoils values (_id, null);
      end if;
    end loop;
  end if;

  return query select
    pirate_count + 1 - id as id, -- 修正海盗分赃的顺序
    amount
  from
    spoils
  order by
    id;
end;
$$ language plpgsql;

spoils.2.sql

create type spoil as (id int, amount int);

create or replace function divide(amount_total int, pirate_count int)
  returns setof spoil as $$
with recursive iters(stage, strategy) as (
  values (1, array[(1, amount_total)::spoil]) -- 初始状态，一个海盗拿全部
union all
  select
    stage + 1, (
    with strategies as ( -- 用嵌套的 with 子句计算到下一个状态的成本
      select
        *,
        case
          when row_number() over () <= (stage + 1) / 2
            then coalesce(amount + 1, 0)
          else 0
        end as cost
      from (
        select
          (unnest(strategy)).*
        order by
          amount nulls first
      ) as t
    )
    select
      case
        when sum(cost) <= amount_total then -- 判断是否能存活
          array_append(array_agg((id, cost)::spoil),
                       (stage + 1, amount_total - sum(cost))::spoil)
        else
          array_append(array_agg((id, amount)::spoil),
                       (stage + 1, null)::spoil)
      end
    from
      strategies
    )
  from
    iters
  where
    stage < pirate_count
)
select
  pirate_count + 1 - id as id,
  amount
from (
  select
    (unnest(strategy)).*
  from
    iters
  where
    stage = pirate_count
  order by
    id desc
) as t
$$ language sql;

spoils.3.sql

create or replace function divide(amount_total int, pirate_count int)
  returns setof record as $$

with recursive spoils(strategy) as (
  values (array[[1, amount_total]]) -- 初始状态，一个海盗拿全部
union all
  select (
    with strategies as ( -- 用嵌套的 with 子句计算到下一个状态的成本
      select
        *,
        case
          when row_number() over () <= (array_length(strategy, 1) + 1) / 2
            then coalesce(amount + 1, 0)
          else 0
        end as cost
      from (
        select
          strategy[i][1] as id,
          strategy[i][2] as amount
        from
          generate_series(array_lower(strategy, 1),
                          array_upper(strategy, 1)) as t(i)
        order by
          amount nulls first
      ) as t
    )
    select
      case
        when sum(cost) <= amount_total then -- 判断是否能存活
          array_cat(array_agg(array[id, cost]),
                    array[[array_length(strategy, 1) + 1,
                           amount_total - sum(cost)]]::int[])
        else
          array_cat(array_agg(array[id, amount]),
                    array[[array_length(strategy, 1) + 1,
                           null]]::int[])
      end
    from
      strategies
  )
  from
    spoils
  where
    array_length(strategy, 1) < pirate_count
)
select
  pirate_count + 1 - strategy[i][1] as id,
  strategy[i][2] as amount
from (
  select
    strategy,
    generate_series(array_lower(strategy, 1),
                    array_upper(strategy, 1)) as i
  from
    spoils
  where
    array_length(strategy, 1) = pirate_count
) as t
order by
  id
$$ language sql;

spoils.4.sql

with recursive spoils(strategy) as (
  values (array[[1, 100]]) -- 初始状态，一个海盗拿全部
union all
  select (
    with strategies as ( -- 用嵌套的 with 子句计算到下一个状态的成本
      select
        *,
        case
          when row_number() over () <= (array_length(strategy, 1) + 1) / 2
            then coalesce(amount + 1, 0)
          else 0
        end as cost
      from (
        select
          strategy[i][1] as id,
          strategy[i][2] as amount
        from
          generate_series(array_lower(strategy, 1),
                          array_upper(strategy, 1)) as t(i)
        order by
          amount nulls first
      ) as t
    )
    select
      case
        when sum(cost) <= 100 then -- 判断是否能存活
          array_cat(array_agg(array[id, cost]),
                    array[[array_length(strategy, 1) + 1,
                           100 - sum(cost)]]::int[])
        else
          array_cat(array_agg(array[id, amount]),
                    array[[array_length(strategy, 1) + 1,
                           null]]::int[])
      end
    from
      strategies
  )
  from
    spoils
  where
    array_length(strategy, 1) < 5
)
select
  5 + 1 - strategy[i][1] as id,
  strategy[i][2] as amount
from (
  select
    strategy,
    generate_series(array_lower(strategy, 1),
                    array_upper(strategy, 1)) as i
  from
    spoils
  where
    array_length(strategy, 1) = 5
) as t
order by
  id;

spoils.5.sql

-- 所有解

with recursive spoils(strategy) as (
  values (array[[1, 100]]) -- 初始状态，一个海盗拿全部[
union all (
with recursive strategies as (
  select
    row_number() over (
      partition by strategy
      order by amount nulls first
    ) as idx,
    *
  from (
    select
      strategy,
      strategy[i][1] as id,
      strategy[i][2] as amount,
      mid
    from (
      select
        strategy,
        (array_length(strategy, 1) + 1) / 2 as mid,
        generate_series(array_lower(strategy, 1),
                        array_upper(strategy, 1)) as i
      from
        spoils
      where
        array_length(strategy, 1) < 5
    ) as t
  ) as t
), medians as (
  select
    strategy,
    idx,
    amount
  from
    strategies
  where
    idx = mid
), groups as (
  select
    strategies.strategy,
    array_agg(array[id, strategies.amount]) filter (where
      strategies.idx <= medians.idx
      and coalesce(strategies.amount, -1) != coalesce(medians.amount, -1)
    ) as prefix,
    array_agg(id) filter (where
      coalesce(strategies.amount, -1) = coalesce(medians.amount, -1)
    ) as middle_indexes,
    max(medians.amount) as middle_amount,
    count(*) filter (where
      strategies.idx <= medians.idx
      and coalesce(strategies.amount, -1) = coalesce(medians.amount, -1)
    ) as middle_size,
    array_agg(array[id, strategies.amount]) filter (where
      strategies.idx > medians.idx
      and coalesce(strategies.amount, -1) != coalesce(medians.amount, -1)
    ) as suffix,
    sum(case
      when strategies.idx <= medians.idx
        then coalesce(strategies.amount + 1, 0)
      else 0
    end) as costs
  from
    strategies
  inner join
    medians
  on
    strategies.strategy = medians.strategy
  group by
    strategies.strategy
), combinations(strategy, prefix, middle, suffix, costs, middle_amount, n, inputs, outputs) as (
  select strategy, prefix, middle_indexes, suffix, costs, middle_amount, middle_size, middle_indexes, '{}'::int[] from groups
union all
  select
    strategy,
    prefix,
    middle,
    suffix,
    costs,
    middle_amount,
    n - 1,
    inputs[idx+1:],
    array_append(outputs, inputs[idx])
  from (
    select
      *,
      generate_series(1, array_length(inputs, 1) - n + 1) as idx
    from
      combinations
    where
      n > 0
  ) as t
)
select
 distinct
  case when costs <= 100
    then array_cat(array_cat(array_cat(array_cat(
      (select array_agg(array[prefix[i][1], coalesce(prefix[i][2] + 1, 0)]) from generate_series(array_lower(prefix, 1), array_upper(prefix, 1)) as t(i)),
      (select array_agg(array[id, coalesce(middle_amount + 1, 0)]) from unnest(outputs) as t(id))),
      (select array_agg(array[id, 0]) from unnest(middle - outputs) as t(id))),
      (select array_agg(array[suffix[i][1], 0]) from generate_series(array_lower(suffix, 1), array_upper(suffix, 1)) as t(i))),
      array[[array_length(strategy, 1) + 1, 100 - costs]]::int[])
    else array_cat(strategy, array[[array_length(strategy, 1) + 1, null]]::int[])
  end
from
  combinations
where
  n = 0
  and middle is not null
  and array_length(middle, 1) > 0)
)
select
  *
from
  spoils
where
  array_length(strategy, 1) = 5;