I have a 50x50 2D dimensional board with empty cells now. I want to fill 20% cells with 0, 30% cells with 1, 30% cells with 2 and 20% cells with 3. How to randomly throw these 4 numbers onto the board with the percentages?
import numpy as np
from numpy import random
dim = 50
map = [[" "for i in range(dim)] for j in range(dim)]
print(map)
One way to get this kind of randomness would be to start with a random permutation of the numbers from 0 to the total number of cells you have minus one.
perm = np.random.permutation(2500)
now you split the permutation according the proportions you want to get and treat the entries of the permutation as the indices of the array.
array = np.empty(2500)
p1 = int(0.2*2500)
p2 = int(0.3*2500)
p3 = int(0.3*2500)
array[perm[range(0, p1)]] = 0
array[perm[range(p1, p1 + p2)]] = 1
array[perm[range(p1 + p2, p3)]] = 2
array[perm[range(p1 + p2 + p3, 2500)]] = 3
array = array.reshape(50, 50)
This way you ensure the proportions for each number.
Since the percentages sum up to 1, you can start with a board of zeros
bsize = 50
board = np.zeros((bsize, bsize))
In this approach the board positions are interpreted as 1D postions, then we need a set of position equivalent to 80% of all positions.
for i, pos in enumerate(np.random.choice(bsize**2, int(0.8*bsize**2), replace=False)):
# the fisrt 30% will be set with 1
if i < int(0.3*bsize**2):
board[pos//bsize][pos%bsize] = 1
# the second 30% (between 30% and 60%) will be set with 2
elif i < int(0.6*bsize**2):
board[pos//bsize][pos%bsize] = 2
# the rest 20% (between 60% and 80%) will be set with 3
else:
board[pos//bsize][pos%bsize] = 3
At the end the last 20% of positions will remain as zeros
As suggested by #alexis in commentaries, this approach could became more simple by using shuffle method from random module:
from random import shuffle
bsize = 50
board = np.zeros((bsize, bsize))
l = list(range(bsize**2))
shuffle(l)
for i, pos in enumerate(l):
# the fisrt 30% will be set with 1
if i < int(0.3*bsize**2):
board[pos//bsize][pos%bsize] = 1
# the second 30% (between 30% and 60%) will be set with 2
elif i < int(0.6*bsize**2):
board[pos//bsize][pos%bsize] = 2
# the rest 20% (between 60% and 80%) will be set with 3
elif i < int(0.8*bsize**2):
board[pos//bsize][pos%bsize] = 3
The last 20% of positions will remain as zeros again.
A different approach (admittedly it's probabilistic so you won't get perfect proportions as the solution proposed by Brad Solomon)
import numpy as np
res = np.random.random((50, 50))
zeros = np.where(res <= 0.2, 0, 0)
ones = np.where(np.logical_and(res <= 0.5, res > 0.2), 1, 0)
twos = np.where(np.logical_and(res <= 0.8, res > 0.5), 2, 0)
threes = np.where(res > 0.8, 3, 0)
final_result = zeros + ones + twos + threes
Running
np.unique(final_result, return_counts=True)
yielded
(array([0, 1, 2, 3]), array([499, 756, 754, 491]))
Here's an approach with np.random.choice to shuffle indices, then filling those indices with repeats of the inserted ints. It will fill the array in the exact proportions that you specify:
import numpy as np
np.random.seed(444)
board = np.zeros(50 * 50, dtype=np.uint8).flatten()
# The "20% cells with 0" can be ignored since that is the default.
#
# This will work as long as the proportions are "clean" ints
# (I.e. mod to 0; 2500 * 0.2 is a clean 500. Otherwise, need to do some rounding.)
rpt = (board.shape[0] * np.array([0.3, 0.3, 0.2])).astype(int)
repl = np.repeat([1, 2, 3], rpt)
idx = np.random.choice(board.shape[0], size=repl.size, replace=False)
board[idx] = repl
board = board.reshape((50, 50))
Resulting frequencies:
>>> np.unique(board, return_counts=True)
(array([0, 1, 2, 3], dtype=uint8), array([500, 750, 750, 500]))
>>> board
array([[1, 3, 2, ..., 3, 2, 2],
[0, 0, 2, ..., 0, 2, 0],
[1, 1, 1, ..., 2, 1, 0],
...,
[1, 1, 2, ..., 2, 2, 2],
[1, 2, 2, ..., 2, 1, 2],
[2, 2, 2, ..., 1, 0, 1]], dtype=uint8)
Approach
Flatten the board. Easier to work with indices when the board is (temporarily) one-dimensional.
rpt is a 1d vector of the number of repeats per int. It gets "zipped" together with [1, 2, 3] to create repl, which is length 2000. (80% of the size of the board; you don't need to worry about the 0s in this example.)
The indices of the flattened array are effectively shuffled (idx), and the length of this shuffled array is constrained to the size of the replacement candidates. Lastly, those indices in the 1d board are filled with the replacements, after which it can be made 2d again.
Related
I have an array of NaNs 10 columns wide and 5 rows long.
I have a 5x3 array of poisson random number generations. This represents 5 runs of each A, B, and C, where each column has a different lambda value for the poisson distribution.
A B C
[1, 1, 2,
1, 2, 2,
2, 1, 4,
1, 2, 3,
0, 1, 2]
Each row represents the number of events. That is, the first row would produce one event of type A, one event of type B, and two events of type C.
I would like to loop through each row and produce a set of uniform random numbers. For A, it would between 1 and 100, for B it would be between 101 and 200, and for C it would be between 201 and 300.
The output of the first row would have four numbers, one number between 1 and 100, one number between 101 and 200, and two numbers between 201 and 300. So a sample output of the first row might be:
[34, 105, 287, 221]
The second output row would have five numbers in it, the third row would have seven, etc. I would like to store it in my array of NaNs by overwriting the NaNs that get replaced in each row. Can anyone please help with this? Thanks!
I've got a rather inefficient/unvectorised method which may or may not be what you're looking for, because one part of your question is unclear to me. Do you want the final array to have rows of different sizes, or to be the same size but padded with nans?
This solution assumes padding with nans, since you talked about the nans being overwritten and didn't mention the extra/unused nans being deleted. I'm also assuming that your ABC thing is structured into a numpy array of size (5,3), and I'm calling the array of nans results_arr.
import numpy as np
from random import randint
# Initializing the arrays
results_arr = np.full((5,10), np.nan)
abc = np.array([[1, 1, 2], [1, 2, 2], [2, 1, 4], [1, 2, 3], [0, 1, 2]])
# Loops through each row in ABC
for row_idx in range(len(abc)):
a, b, c = abc[row_idx]
# Here, I'm getting a number in the specified uniform distribution as many times as is specified in the A column. The other 2 loops do the same for the B and C columns.
for i in range(0, a):
results_arr[row_idx, i] = randint(1, 100)
for j in range(a, a+b):
results_arr[row_idx, j] = randint(101, 200)
for k in range(a+b, a+b+c):
results_arr[row_idx, k] = randint(201, 300)
Hope that helps!
P.S. Here's a solution with uneven rows. The result is stored in a list of lists because numpy doesn't support ragged arrays (i.e. rows of different lengths).
import numpy as np
from random import randint
# Initializations
results_arr = []
abc = np.array([[1, 1, 2], [1, 2, 2], [2, 1, 4], [1, 2, 3], [0, 1, 2]])
# Same code logic as before, just storing the results differently
for row_idx in range(len(abc)):
a, b, c = abc[row_idx]
results_this_row = []
for i in range(0, a):
results_this_row.append(randint(1, 100))
for j in range(a, a+b):
results_this_row.append(randint(101, 200))
for k in range(a+b, a+b+c):
results_this_row.append(randint(201, 300))
results_arr.append(results_this_row)
I hope these two solutions cover what you're looking for!
I would like to replace the N smallest elements in each row for 0, and that the resulting array would respect the same order and shape of the original array.
Specifically, if the original numpy array is:
import numpy as np
x = np.array([[0,50,20],[2,0,10],[1,1,0]])
And N = 2, I would like for the result to be the following:
x = np.array([[0,50,0],[0,0,10],[0,1,0]])
I tried the following, but in the last row it replaces 3 elements instead of 2 (because it replaces both 1s and not only one)
import numpy as np
N = 2
x = np.array([[0,50,20],[2,0,10],[1,1,0]])
x_sorted = np.sort(x , axis = 1)
x_sorted[:,N:] = 0
replace = x_sorted.copy()
final = np.where(np.isin(x,replace),0,x)
Note that this is small example and I would like that it works for a much bigger matrix.
Thanks for your time!
One way using numpy.argsort:
N = 2
x[x.argsort().argsort() < N] = 0
Output:
array([[ 0, 50, 0],
[ 0, 0, 10],
[ 0, 1, 0]])
Use numpy.argpartition to find the index of N smallest elements, and then use the index to replace values:
N = 2
idy = np.argpartition(x, N, axis=1)[:, :N]
x[np.arange(len(x))[:,None], idy] = 0
x
array([[ 0, 50, 0],
[ 0, 0, 10],
[ 1, 0, 0]])
Notice if there are ties, it could be undetermined which values get replaced depending on the algorithm used.
Let's consider the distance d(a, b) = number of digits which are pairwise different in a and b, e.g.:
d(1003000000, 1000090000) = 2 # the 4th and 6th digits don't match
(we only work with 10-digit numbers) and this list:
L = [2678888873,
2678878873, # distance 1 from L[0]
1000000000,
1000040000, # distance 1 from L[2]
1000300000, # distance 1 from L[2], distance 2 from L[3]
1000300009, # distance 1 from L[4], distance 2 from L[2]
]
I would like to find the minimal number of points P such that each integer in the list is at a distance <= 1 of a point in P.
Here I think this number is 3: every number in the list is at distance <= 1 of 2678888873, 1000000000, or 1000300009.
I imagine an O(n^2) algorithm is possible by first computing a distance matrix i.e. M[i, j] = d(L[i], L[j]).
Is there a better way to do this, especially using Numpy? (maybe there's a built-in algorithm in Numpy/Scipy?)
PS: If we see these 10-digit integers as strings, we're close to finding a minimal number of clusters in a list of many words with a Levenshtein distance.
PS2: I know realize this distance has a name on strings: Hamming distance.
Let's see what we know from a the distance metric. Given a number P (not necessarily in L), if two members of L are within distance 1 of P, they each share 9 digits with P, but not necessarily the same ones, so they are only guaranteed to share 8 digits with each other. So any two numbers that have distance 2 are guaranteed to two unique Ps that are distance 1 from each of them (and distance 2 from each other as well). You can use this information to reduce the amount of brute-force effort required to optimize the selection of P.
Let's say you have a distance matrix. You can immediately discard rows (or columns) that don't have entries less than 3: they are their own cluster automatically. For the remaining entries that are equal to 2, construct a list of possible P values. Find the number of elements of L that are within 1 of each element of P (another distance matrix). Sort P by the number of neighbors, and select. You will need to update the matrix at each iteration as you remove members with maximal neighbors to avoid inefficient grouping due to overlap (members of L that are near multiple members of P).
You can compute a distance matrix for L in numpy by first converting it to a 2D array of digits:
L = np.array([2678888873, 2678878873, 1000000000, 1000040000, 1000300000, 1000300009])
z = 10 # Number of digits
n = len(L) # Number of numbers
dec = 10**np.arange(z).reshape(-1, 1).astype(np.int64)
digits = (L // dec) % 10
digits is now a 10xN array:
array([[3, 3, 0, 0, 0, 9],
[7, 7, 0, 0, 0, 0],
[8, 8, 0, 0, 0, 0],
[8, 8, 0, 0, 0, 0],
[8, 7, 0, 4, 0, 0],
[8, 8, 0, 0, 3, 3],
[8, 8, 0, 0, 0, 0],
[7, 7, 0, 0, 0, 0],
[6, 6, 0, 0, 0, 0],
[2, 2, 1, 1, 1, 1]], dtype=int64)
You can compute the distance between digits and itself, or digits and any other 10xM array using != and sum along the right axis:
distance = (digits[:, None, :] != digits[..., None]).sum(axis=0)
The result:
array([[ 0, 1, 10, 10, 10, 10],
[ 1, 0, 10, 10, 10, 10],
[10, 10, 0, 1, 1, 2],
[10, 10, 1, 0, 2, 3],
[10, 10, 1, 2, 0, 1],
[10, 10, 2, 3, 1, 0]])
We are only concerned with the upper (or lower) triangle of that matrix, so we can immediately mask out the other triangle:
distance[np.tril_indices(n)] = z + 1
Find all candidate values of P: all elements of L, but also all pairs between elements that have distance 2:
# Find indices of pairs that differ by 2
indices = np.nonzero(distance == 2)
# Extract those numbers as 10xKx2 array
d = digits[:, np.stack(indices, axis=1)]
# Compute where the difference is nonzero (Kx2)
locs = np.diff(d, axis=2).astype(bool).squeeze()
# Find the index of the first digit to replace (K)
s = np.argmax(locs, axis=0)
The extra values of P are constructed from each half of d, with the digits represented by k replaced from the other half:
P0 = digits[:, indices[0]]
P1 = digits[:, indices[1]]
k = np.arange(s.size)
tmp = P0[s, k]
P0[s, k] = P1[s, k]
P1[s, k] = tmp
Pextra = np.unique(np.concatenate((P0, P1), axis=1)
So now you can compute the total set of possibilities for P:
P = np.concatenate((digits, Pextra), axis=1)
distance2 = (P[:, None, :] != digits[..., None]).sum(axis=0)
You can discard any elements of Pextra that match with elements of digits based on the distance:
mask = np.concatenate((np.ones(n, bool), distance2[:, n:].all(axis=0)))
P = P[:, mask]
distance2 = distance2[:, mask]
Now you can iteratively distance P with L, and select the best values of P, removing any values that have been selected from the distance matrix. A greedy selection from P will not necessarily be optimal, since an alternative combination may require fewer elements due to overlaps, but that is a matter for a simple (but somewhat expensive) graph traversal algorithm. The following snippet just shows a simple greedy selection, which will work fine for your toy example:
distMask = distance2 <= 1
quality = distMask.sum(axis=0)
clusters = []
accounted = 0
while accounted < n:
# Get the cluster location
best = np.argmax(quality)
# Get the cluster number
clusters.append(P[:, best].dot(dec).item())
# Remove numbers in cluser from consideration
accounted += quality[best]
quality -= distMask[distMask[:, best], :].sum(axis=0)
The last couple of steps can be optimized using sets and graphs, but this shows a starting point for a valid approach. This is going to be slow for large data, but probably not prohibitively so. Do some benchmarks to decide how much time you want to spend optimizing vs just running the algorithm.
I'm trying to convert my MATLAB code to python but I'm having some issues. This code is supposed to segment letters from a picture.
Here's the whole code in MATLAB:
he = imread('r.jpg');
imshow(he);
%C = makecform(type) creates the color transformation structure C that defines the color space conversion specified by type.
cform = makecform('srgb2lab');
%To perform the transformation, pass the color transformation structure as an argument to the applycform function.
lab_he = applycform(he,cform);
%convert to double precision
ab = double(lab_he(:,:,2:3));
%size of dimension in 2D array
nrows = size(ab,1);
ncols = size(ab,2);
%B = reshape(A,sz1,...,szN) reshapes A into a sz1-by-...-by-szN array where
%sz1,...,szN indicates the size of each dimension. You can specify a single
% dimension size of [] to have the dimension size automatically calculated,
% such that the number of elements in B matches the number of elements in A.
% For example, if A is a 10-by-10 matrix, then reshape(A,2,2,[]) reshapes
% the 100 elements of A into a 2-by-2-by-25 array.
ab = reshape(ab,nrows*ncols,2);
% repeat the clustering 3 times to avoid local minima
nColors = 3;
[cluster_idx, cluster_center] = kmeans(ab,nColors,'distance','sqEuclidean', ...
'Replicates',3);
pixel_labels = reshape(cluster_idx,nrows,ncols);
imshow(pixel_labels,[]), title('image labeled by cluster index');
segmented_images = cell(1,3);
rgb_label = repmat(pixel_labels,[1 1 3]);
for k = 1:nColors
color = he;
color(rgb_label ~= k) = 0;
segmented_images{k} = color;
end
figure,imshow(segmented_images{1}), title('objects in cluster 1');
figure,imshow(segmented_images{2}), title('objects in cluster 2');
figure,imshow(segmented_images{3}), title('objects in cluster 3');
mean_cluster_value = mean(cluster_center,2);
[tmp, idx] = sort(mean_cluster_value);
blue_cluster_num = idx(1);
L = lab_he(:,:,1);
blue_idx = find(pixel_labels == blue_cluster_num);
L_blue = L(blue_idx);
is_light_blue = im2bw(L_blue,graythresh(L_blue));
% target_labels = repmat(uint8(0),[nrows ncols]);
% target_labels(blue_idx(is_light_blue==false)) = 1;
% target_labels = repmat(target_labels,[1 1 3]);
% blue_target = he;
% blue_target(target_labels ~= 1) = 0;
% figure,imshow(blue_target), title('blue');
Here's what I have in Python so far:
import cv2
import numpy as np
from matplotlib import pyplot as plt
import sys
img = cv2.imread('r.jpg',1)
print "original image: ", img.shape
cv2.imshow('BGR', img)
img1 = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
img2 = cv2.cvtColor(img1, cv2.COLOR_RGB2LAB)
cv2.imshow('RGB', img1)
cv2.imshow('LAB', img2) #differs from the LAB color space in MATLAB (need to patch maybe?)
print "LAB converted image: ", img2.shape
print "LAB converted image dimension", img2.ndim #says the image is a 3 dimensional array
img2a = img2.shape[2][1:2]
print "Slicing the LAB converted image", img2a
#we need to convert that to double precision
print img2.dtype
img2a = img2.astype(np.uint64) #convert to double precision
print img2a.dtype
#print img2a
row = img2a.shape[0] #returns number of rows of img2a
column = img2a.shape[1] #returns number of columns of img2a
print "row: ", row #matches the MATLAB version
print "column: ", column #matchees the MATLAB version
rowcol = row * column
k = cv2.waitKey(0)
if k == 27: # wait for ESC key to exit
cv2.destroyAllWindows()
elif k == ord('s'): # wait for 's' key to save and exit
cv2.imwrite('final image',final_image)
cv2.destroyAllWindows()
Now the part i'm currently stuck in is that here in Matlab code, I have lab_he(:,:,2:3) which means all the rows and all the columns from depth 2 and 3. However when I try to replicate that in Python img2a = img2.shape[2][1:2] but it doesn't work or makes sense. Please help.
In Octave/MATLAB
octave:29> x=reshape(1:(2*3*4),3,2,4);
octave:30> x(:,:,2:3)
ans =
ans(:,:,1) =
7 10
8 11
9 12
ans(:,:,2) =
13 16
14 17
15 18
octave:31> size(x(:,:,2:3))
ans =
3 2 2
octave:33> x(:,:,2:3)(2,2,:)
ans(:,:,1) = 11
ans(:,:,2) = 17
In numpy:
In [13]: x=np.arange(1,1+2*3*4).reshape(3,2,4,order='F')
In [14]: x[:,:,1:3]
Out[14]:
array([[[ 7, 13],
[10, 16]],
[[ 8, 14],
[11, 17]],
[[ 9, 15],
[12, 18]]])
In [15]: _.shape
Out[15]: (3, 2, 2)
In [17]: x[:,:,1:3][1,1,:]
Out[17]: array([11, 17])
Or with numpy normal 'C' ordering, and indexing on the 1st dimension
In [18]: y=np.arange(1,1+2*3*4).reshape(4,2,3)
In [19]: y[1:3,:,:]
Out[19]:
array([[[ 7, 8, 9],
[10, 11, 12]],
[[13, 14, 15],
[16, 17, 18]]])
In [20]: y[1:3,:,:][:,1,1]
Out[20]: array([11, 17])
Same indexing ideas, though matching numbers and shapes requires some care, not only with the 0 v 1 index base. A 3d array is displayed in a different arangement. Octave divides it into blocks on the last index (its primary iterator), numpy iterates on the first index.
In 3d it makes more sense to talk about first, 2nd, 3rd dimensions rather than row,col,depth. In 4d you run out of names. :)
I had to reshape array at specific depth, and I programmed a little recursive function to do so:
def recursive_array_cutting(tab, depth, i, min, max):
if(i==depth):
tab = tab[min:max]
return tab
temp_list = []
nb_subtab = len(tab)
for index in range(nb_subtab):
temp_list.append(recursive_array_cutting(tab[index], depth, i+1, min, max))
return np.asanyarray(temp_list)
It allow to get all array/values between the min and the max of a specific depth, for instance, if you have a (3, 4) tab = [[0, 1, 2, 3], [0, 1, 2, 3], [0, 1, 2, 3]] and only want the last two values of the deepest array, you call like this : tab = recursive_array_cutting(tab, 1, 0, 0, 2) to get the output : [[0 1][0 1][0 1]].
If you have a more complexe array like this tab = [[[0, 1, 2, 3], [1, 1, 2, 3], [2, 1, 2, 3]], [[0, 1, 2, 3], [1, 1, 2, 3], [2, 1, 2, 3]], [[0, 1, 2, 3], [1, 1, 2, 3], [2, 1, 2, 3]]] (3, 3, 4) and want a (3, 2, 4) array, you can call like this : tab = recursive_array_cutting(tab, 1, 0, 0, 2) to get this output, and get rid of the last dimension in depth 1.
Function like this surely exist in numpy, but I did not found it.
Given the following array:
complete_matrix = numpy.array([
[0, 1, 2, 4],
[1, 0, 3, 5],
[2, 3, 0, 6],
[4, 5, 6, 0]])
I would like to identify the row with the highest average, excluding the diagonal zeros.
So, in this case, I would be able to identify complete_matrix[:,3] as being the row with the highest average.
Note that the presence of the zeros doesn't affect which row has the highest mean because all rows have the same number of elements. Therefore, we just take the mean of each row, and then ask for the index of the largest element.
#Take the mean along the 1st index, ie collapse into a Nx1 array of means
means = np.mean(complete_matrix, 1)
#Now just get the index of the largest mean
idx = np.argmax(means)
idx is now the index of the row with the highest mean!
You don't need to worry about the 0s, they shouldn't effect how the averages compare since there will presumably be one in each row. Hence, you can do something like this to get the index of the row with the highest average:
>>> import numpy as np
>>> complete_matrix = np.array([
... [0, 1, 2, 4],
... [1, 0, 3, 5],
... [2, 3, 0, 6],
... [4, 5, 6, 0]])
>>> np.argmax(np.mean(complete_matrix, axis=1))
3
Reference:
numpy.mean
numpy.argmax
As pointed out by a lot of people, presence of zeros isn't an issue as long as you have the same number of zeros in each column. Just in case your intention was to ignore all the zeros, preventing them from participating in the average computation, you could use weights to suppress the contribution of the zeros. The following solution assigns 0 weight to zero entries, 1 otherwise:
numpy.argmax(numpy.average(complete_matrix,axis=0, weights=complete_matrix!=0))
You can always create a weight matrix where the weight is 0 for diagonal entries, and 1 otherwise.
You will see that this answer actually would fit better to your other question that was marked as duplicated to this one (and don't know why because it is not the same question...)
The presence of zeros can indeed affect the columns' or rows' average, for instance:
a = np.array([[ 0, 1, 0.9, 1],
[0.9, 0, 1, 1],
[ 1, 1, 0, 0.5]])
Without eliminating the diagonals, it would tell that the column 3 has the highest average, but eliminating the diagonals the highest average belongs to column 1 and now column 3 has the least average of all columns!
You can correct the calculated mean using the lcm (least common multiple) of the number of lines with and without the diagonals, by guaranteeing that where a diagonal element does not exist the correction is not applied:
correction = column_sum/lcm(len(column), len(column)-1)
new_mean = mean + correction
I copied the algorithm for lcm from this answer and proposed a solution for your case:
import numpy as np
def gcd(a, b):
"""Return greatest common divisor using Euclid's Algorithm."""
while b:
a, b = b, a % b
return a
def lcm(a, b):
"""Return lowest common multiple."""
return a * b // gcd(a, b)
def mymean(a):
if len(a.diagonal()) < a.shape[1]:
tmp = np.hstack((a.diagonal()*0+1,0))
else:
tmp = a.diagonal()*0+1
return np.mean(a, axis=0) + np.sum(a,axis=0)*tmp/lcm(a.shape[0],a.shape[0]-1)
Testing with the a given above:
mymean(a)
#array([ 0.95 , 1. , 0.95 , 0.83333333])
With another example:
b = np.array([[ 0, 1, 0.9, 0],
[0.9, 0, 1, 1],
[ 1, 1, 0, 0.5],
[0.9, 0.2, 1, 0],
[ 1, 1, 0.7, 0.5]])
mymean(b)
#array([ 0.95, 0.8 , 0.9 , 0.5 ])
With the corrected average you just use np.argmax() to get the column index with the highest average. Similarly, np.argmin() to get the index of the column with the least average:
np.argmin(mymean(a))