simple setindex

Author: mcabbott

src/array.jl

@@ -80,6 +80,16 @@ end
 Base.getindex(x::OneHotArray, ::Colon) = BitVector(reshape(x, :))
 Base.getindex(x::OneHotArray{<:Any, N}, ::Colon, ::Vararg{Colon, N}) where N = x
 
+@inline function Base.setindex!(x::OneHotArray{<:Any, N}, v, i::Integer, I::Vararg{Integer, N}) where N
+  @boundscheck checkbounds(x, i, I...)
+  if Bool(v)
+    @inbounds x.indices[I...] = i
+  elseif x.indices[I...] == i
+    @inbounds x.indices[I...] = 0
+  end
+  x
+end
+
 function Base.showarg(io::IO, x::OneHotArray, toplevel)
   print(io, ndims(x) == 1 ? "OneHotVector(" : ndims(x) == 2 ? "OneHotMatrix(" : "OneHotArray(")
   Base.showarg(io, x.indices, false)