Elixir/lib/data_structures/seg_tree.ex at 41cea21cd76631214d53dda87f822e7e7a21cc98 · TheAlgorithms/Elixir · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
defmodule Algorithms.DataStructures.SegTree do
  @moduledoc """
  Generic Segment Tree implementation in Elixir.

  Usage example:

      op = fn a, b -> min(a, b) end
      e = fn -> :infinity end

      seg = Algorithms.DataStructures.SegTree.new([5, 3, 7, 9, 6], op, e)
      Algorithms.DataStructures.SegTree.prod(seg, 1, 4)   # => 3
      seg = Algorithms.DataStructures.SegTree.set(seg, 2, 1)
      Algorithms.DataStructures.SegTree.prod(seg, 1, 4)   # => 1
  """

  defstruct size: 0, n: 0, data: nil, op: nil, e: nil

  @type t :: %__MODULE__{
          size: non_neg_integer(),
          n: non_neg_integer(),
          data: :array.array(),
          op: (any(), any() -> any()),
          e: (() -> any())
        }

  @spec new(list(any()), (any(), any() -> any()), (() -> any())) :: t()
  def new(list, op, e) do
    n = length(list)
    size = bit_ceil(n)
    data = :array.new(size * 2, default: e.())

    data =
      Enum.with_index(list)
      |> Enum.reduce(data, fn {v, i}, acc ->
        :array.set(size + i, v, acc)
      end)

    data = build_tree(data, size, op)

    %__MODULE__{size: size, n: n, data: data, op: op, e: e}
  end

  defp build_tree(data, size, op) do
    Enum.reduce(Enum.reverse(1..(size - 1)), data, fn i, acc ->
      left = :array.get(2 * i, acc)
      right = :array.get(2 * i + 1, acc)
      :array.set(i, op.(left, right), acc)
    end)
  end

  @spec set(t(), non_neg_integer(), any()) :: t()
  def set(%__MODULE__{size: size, data: data, op: op} = st, p, x) do
    data = :array.set(size + p, x, data)
    data = update_up(size + p, data, op)
    %__MODULE__{st | data: data}
  end

  defp update_up(1, data, _op), do: data

  defp update_up(k, data, op) do
    parent = div(k, 2)
    left = :array.get(2 * parent, data)
    right = :array.get(2 * parent + 1, data)
    data = :array.set(parent, op.(left, right), data)
    update_up(parent, data, op)
  end

  @spec get(t(), non_neg_integer()) :: any()
  def get(%__MODULE__{size: size, data: data}, p) do
    :array.get(size + p, data)
  end

  @spec prod(t(), non_neg_integer(), non_neg_integer()) :: any()
  def prod(%__MODULE__{size: size, data: data, op: op, e: e}, l, r) do
    do_prod(l + size, r + size, e.(), e.(), data, op)
  end

  defp do_prod(l, r, sml, smr, data, op) do
    do_prod_loop(l, r, sml, smr, data, op)
  end

  defp do_prod_loop(l, r, sml, smr, data, op) when l < r do
    sml =
      if rem(l, 2) == 1 do
        op.(sml, :array.get(l, data))
      else
        sml
      end

    l = if rem(l, 2) == 1, do: l + 1, else: l

    smr =
      if rem(r, 2) == 1 do
        op.(:array.get(r - 1, data), smr)
      else
        smr
      end

    r = if rem(r, 2) == 1, do: r - 1, else: r

    do_prod_loop(div(l, 2), div(r, 2), sml, smr, data, op)
  end

  defp do_prod_loop(_, _, sml, smr, _, op), do: op.(sml, smr)

  @spec all_prod(t()) :: any()
  def all_prod(%__MODULE__{data: data}), do: :array.get(1, data)

  @spec bit_ceil(non_neg_integer()) :: non_neg_integer()
  defp bit_ceil(0), do: 1

  defp bit_ceil(n) when n > 0 do
    pow = :math.ceil(:math.log2(n))
    trunc(:math.pow(2, pow))
  end
end