SparseMatrix

SparseMatrix.hx

@:allow(SparseMatrixIterator)
@:allow(SparseMatrixKeyValueIterator)
class SparseMatrix<T> {
    final _rows:Array<Int>;
    final _cols:Array<Int>;
    final _values:Array<T>;

    public function new() {
        _rows = new Array();
        _cols = new Array();
        _values = new Array();
    }

    public function add(x:Int, y:Int, value:T):Void {
        _rows.push(x);
        _cols.push(y);
        _values.push(value);
    }

    public function get(x:Int, y:Int):Null<T> {
        for (i in 0..._rows.length)
            if (_rows[i] == x && _cols[i] == y)
                return _values[i];
        return null;
    }

    public function remove(x:Int, y:Int):Null<T> {
        for (i in 0..._rows.length) {
            if (_rows[i] == x && _cols[i] == y) {
                _rows.splice(i, 1);
                _cols.splice(i, 1);
                return _values.splice(i, 1)[0];
            }
        }
        return null;
    }

    public function iterator()
        return new SparseMatrixIterator(this);

    public function keyValueIterator()
        return new SparseMatrixKeyValueIterator(this);
}

class SparseMatrixIterator<T> {
    var _matrix:SparseMatrix<T>;
    var _i:Int;

    public inline function new(matrix:SparseMatrix<T>) {
        _matrix = matrix;
        _i = 0;
    }

    public inline function hasNext():Bool {
        return _i < _matrix._values.length;
    }

    public inline function next() {
        return _matrix._values[_i++];
    }
}

class SparseMatrixKeyValueIterator<K, V> {
    var _matrix:SparseMatrix<V>;
    var _i:Int;

    public inline function new(matrix:SparseMatrix<V>) {
        _matrix = matrix;
        _i = 0;
    }

    public inline function hasNext():Bool {
        return _i < _matrix._values.length;
    }

    public inline function next():{key:{x:Int, y:Int}, value:V} {
        return {
            key: { x: _matrix._rows[_i], y: _matrix._cols[_i] },
            value: _matrix._values[_i++],
        };
    }
}