How can I modify given Matlab code for 2D kmeans to 3D

I am given this two pieces of matlab code

% --- initialize Matlab
clc;
disp('-----');
clear variables;
close all force;
dbstop if error;
dbstop if warning;

% --- initialize the random generator 
rng('default');

% --- select a file (image) to load the 3D data points (pixels)
Im=imread('toy5.jpg');

% --- define k
k=4;

% --- get image data in a Nx3 array (r,g,b triplets)
X=double(reshape(Im,[],3));

% --- get the overall number of pixels
N=size(X,1);

C=[];
while size(unique(C,'rows'),1)~=k
    % --- find k random unique pixels (i.e. with different rgb values)
    Idx=round(rand(k,1)*N);

    % --- set these pixels as initial cluster centers
    C=X(Idx,:);
end

% --- show the pixels (every 500-th) as points in the RGB space
scatter3(X(1:200:end,1),X(1:200:end,2),X(1:200:end,3),'k.');
hold on;

% --- show the cluster centers with k different colors
scatter3(C(:,1),C(:,2),C(:,3),100,(1:k)','o','filled','MarkerEdgeColor','k');

% --- set drawing parameters
axis equal;
box on;
grid off;
view([-17 30]);
xlabel('Red','color','r');
ylabel('Green','color','g');
zlabel('Blue','color','b');
colormap(hsv(k));

and

    % --- initialize Matlab
disp('-----');
clear variables;
close all force;
dbstop if error;
dbstop if warning;

% --- select a file to load the 2D data points
fprintf(2,'Enter the DATA POINTS filename:\n');
[FileName1,PathName1]= uigetfile({'*.txt';'*.*'},'Open');
A=load([PathName1 FileName1]);

% --- select a file to load the 2D initial clulster centers
fprintf(2,'Enter the CLUSTER CENTERS filename:\n');
[FileName2,PathName2]= uigetfile({'*.txt';'*.*'},'Open');
C=load([PathName2 FileName2]);

% --- get the number of data points
N=size(A,1);

% --- get the number of clusters K
K=size(C,1);

% --- define the cluster colors in the figure according to number K
Colors=hsv(max(3,K));
colormap(Colors);

% --- plot the data points and the original centers
hold on;
scatter(A(:,1),A(:,2),30,'ko','filled');
scatter(C(:,1),C(:,2),100,(1:K)','o','filled','MarkerEdgeColor','k');

% --- set drawing parameters
set(gca,'ticklength',[0 0]);
box on;
axis image;
axis([0 max([A(:,1);C(:,1)])+20 0 max([A(:,2);C(:,2)])+20]);

% --- wait for a key to be pressed
title('Press a key to start...','FontSize',14,'Color','r');
pause;

Itteration=0;
while Itteration<10
    % --- increase the itterations counter
    Itteration=Itteration+1;

    % --- find the distance of all N points from the K clusters.
    % --- the resulting array Dist has dimensions [NxK].
    % --- for example Dist(7,2) is the distance of the 7th data point from the 2nd center.
    for i=1:K
        Dist(:,i)=sqrt(sum((A-repmat(C(i,:),N,1)).^2,2));
    end

    % --- for each one of the N data points, find the minimum of the K distances (MinValue)
    % --- and the corresponding cluster (T)
    [MinValue,T]=min(Dist,[],2);

    % --- update each cluster's center by the points that are assigned to it
    for i=1:K
        if sum(T==i)>0
            C(i,:)=(mean(A(T==i,:),1));
        end
    end

    % --- color the data points according to the cluster they belong
    % --- and show the updated centers position
    cla;
    scatter(A(:,1),A(:,2),30,T,'o','filled');
    scatter(C(:,1),C(:,2),100,(1:K)','o','filled','MarkerEdgeColor','k');

    disp(C); % show the k centers in the Command Window
    title(['Itteration ' num2str(Itteration) ' - Press a key to continue...'],'FontSize',14,'Color','r');
    pause;
end

I am supposed to combine this code to produce code for the kmeans algorithm in 3D. I did this:

% --- initialize Matlab
clc; disp('-----'); 
clear variables; close all force; dbstop if error; dbstop if warning;

% --- initialize the random generator 
rng('default');
% --- select a file (image) to load the 3D data points (pixels)
Im=imread('toy5.jpg');
% --- define k
k=4;
% --- get image data in a Nx3 array (r,g,b triplets)
X=double(reshape(Im,[],3));
% --- get the overall number of pixels
N=size(X,1);
C=[];
while size(unique(C,'rows'),1)~=k
    % --- find k random unique pixels (i.e. with different rgb values)
    Idx=round(rand(k,1)*N);   
    % --- set these pixels as initial cluster centers
    C=X(Idx,:);
end

% --- show the pixels (every 500-th) as points in the RGB space
scatter3(X(1:200:end,1),X(1:200:end,2),X(1:200:end,3),'k.');
hold on;

% --- show the cluster centers with k different colors
scatter3(C(:,1),C(:,2),C(:,3),100,(1:k)','o','filled','MarkerEdgeColor','k');

% --- set drawing parameters
axis equal; box on; grid off; view([-17 30]);
% xlabel('Red','color','r'); ylabel('Green','color','g');
% zlabel('Blue','color','b'); colormap(hsv(k));


% --- wait for a key to be pressed
title('Press a key to start...','FontSize',14,'Color','r');
pause;

Itteration=0;
LastC=zeros(size(C));
while true
    % --- increase the itterations counter
    Itteration=Itteration+1;

    % --- find the distance of all N points from the K clusters.
    % --- the resulting array Dist has dimensions [NxK].
    % --- for example Dist(7,2) is the distance of the 7th data point from the 2nd center.
    for i=1:k
        Dist(:,i)=sqrt(sum((X-repmat(C(i,:,:),N,1)).^2,2));
    end

    % --- for each one of the N data points, find the minimum of the K distances (MinValue)
    % --- and the corresponding cluster (T)
    [MinValue,T]=min(Dist,[],2);

    % --- update each cluster's center by the points that are assigned to it
    for i=1:k
        if sum(T==i)>0
            C(i,:)=(mean(X(T==i,:,:),1));
        end
    end

    % --- and show the updated centers position
    cla;

    % --- show the pixels (every 50-th) as points in the RGB space
    scatter3(X(1:50:end,1),X(1:50:end,2),X(1:50:end,3),'k.');

    % --- show the cluster centers with k different colors
    scatter3(C(:,1),C(:,2),C(:,3),100,(1:k)','o','filled','MarkerEdgeColor','k');

    % --- if the last and the current centers positions are not the same then
    % --- show the itterations number in the title continue...
    if any(any(LastC-C~=0))
        LastC=C;
        disp(C); % show the k centers in the Command Window
        title(['Itteration ' num2str(Itteration) ],'FontSize',14,'Color','r');
        drawnow; %pause;
    else
        % --- otherwise show the final itterations number in the title and finish
        title(['Itteration ' num2str(Itteration) ' - Finished!'],'FontSize',14,'Color','b');
        break;
    end
end

% open a new figure, set the same drawing parameter
figure;
 % --- show the pixels (every 50-th) COLORED as their class
scatter3(X(1:50:end,1),X(1:50:end,2),X(1:50:end,3),5,T(1:50:end),'o','filled');
hold on;
scatter3(C(:,1),C(:,2),C(:,3),100,(1:k)','o','filled','MarkerEdgeColor','k');
axis equal; box on; grid off; view([-17 30]);
% xlabel('Red','color','r'); ylabel('Green','color','g');
% zlabel('Blue','color','b'); colormap(hsv(k));
title(['Pixels colored as the class that they belong'],'FontSize',14,'Color','b');

newT= reshape(T,300,300);
figure;
image(newT);
colormap(C/255);
axis image; 

and it seems it give correct results, but when I proceed to the next steps of the assignment I am getting wrong values.

I have to modify the third piece of code so that at each Iteration we calculate the Residual sum of squares (a function JRoc = J(iter)-J(iter-1)/J(iter) is given).

But the results are not the expected ones.

This is the code of the next step

% --- initialize Matlab
clc; disp('-----'); 
clear variables; close all force; dbstop if error; dbstop if warning;

% --- initialize the random generator 
rng('default');
% --- select a file (image) to load the 3D data points (pixels)
Im=imread('buildings1.jpg');
% --- define k
k=8;
% --- get image data in a Nx3 array (r,g,b triplets)
X=double(reshape(Im,[],3));
% --- get the overall number of pixels
N=size(X,1);
C=[];
while size(unique(C,'rows'),1)~=k
    % --- find k random unique pixels (i.e. with different rgb values)
    Idx=round(rand(k,1)*N);   
    % --- set these pixels as initial cluster centers
    C=X(Idx,:);
end

% --- show the pixels (every 500-th) as points in the RGB space
scatter3(X(1:300:end,1),X(1:300:end,2),X(1:300:end,3),'k.');
hold on;

% --- show the cluster centers with k different colors
scatter3(C(:,1),C(:,2),C(:,3),100,(1:k)','o','filled','MarkerEdgeColor','k');

% --- set drawing parameters
axis equal; box on; grid off; view([-17 30]);

Itteration=0;
MaxItteration=70; 
Minjsum=2.3*10^8; jsum=10^11; 
Minjroc=0.0001; lastjroc=-14; 
LastC=zeros(size(C));
while and(Itteration<=MaxItteration, jsum>Minjsum); % --- Katofli J &MaxI
%while and(Itteration<=MaxItteration, lastjroc<Minjroc);
%while true

    % --- increase the itterations counter
    Itteration=Itteration+1;

    % --- find the distance of all N points from the K clusters.
    % --- the resulting array Dist has dimensions [NxK].
    % --- for example Dist(7,2) is the distance of the 7th data point from the 2nd center.
    for i=1:k
        Dist(:,i)=sqrt(sum((X-repmat(C(i,:,:),N,1)).^2,2));
    end

    % --- for each one of the N data points, find the minimum of the K distances (MinValue)
    % --- and the corresponding cluster (T)
    [MinValue,T]=min(Dist,[],2);

    % --- Ahtroisma tetragonon apostaseon apo kathe klassi
    jsum=0;
    for j=1:k % gia kathe klassi
            jsum=jsum+sum((Dist(:,j)).^2);
    end
    jsum_iter(Itteration)=jsum;

    % --- Jroc reduction rythm
    if Itteration>1 
        jroc(Itteration)=(jsum_iter(Itteration)-jsum_iter(Itteration-1))/jsum_iter(Itteration-1)*100;
        lastjroc=jroc(Itteration);
    end

    % --- update each cluster's center by the points that are assigned to it
    for i=1:k
        if sum(T==i)>0
            C(i,:)=(mean(X(T==i,:,:),1));
        end
    end

    % --- and show the updated centers position
    cla;

    % --- show the pixels (every 350-th) as points in the RGB space
    scatter3(X(1:400:end,1),X(1:400:end,2),X(1:400:end,3),'k.');

    % --- show the cluster centers with k different colors
    scatter3(C(:,1),C(:,2),C(:,3),100,(1:k)','o','filled','MarkerEdgeColor','k');

    % --- if the last and the current centers positions are not the same then
    % --- show the itterations number in the title continue...
    if any(any(LastC-C~=0))
        LastC=C;
        disp(C); % show the k centers in the Command Window
        title(['Itteration ' num2str(Itteration) ],'FontSize',14,'Color','r');
        drawnow; %pause;
    else
        % --- otherwise show the final itterations number in the title and finish
        title(['Itteration ' num2str(Itteration) ' - Finished!'],'FontSize',14,'Color','b');
        break;
    end
end

% open a new figure, set the same drawing parameter
figure;
 % --- show the pixels (every 50-th) COLORED as their class
scatter3(X(1:300:end,1),X(1:300:end,2),X(1:300:end,3),5,T(1:300:end),'o','filled');
hold on;
scatter3(C(:,1),C(:,2),C(:,3),100,(1:k)','o','filled','MarkerEdgeColor','k');
axis equal; box on; grid off; view([-17 30]);
% xlabel('Red','color','r'); ylabel('Green','color','g');
% zlabel('Blue','color','b'); colormap(hsv(k));
title(['Pixels colored as the class that they belong'],'FontSize',14,'Color','b');

newT= reshape(T,616,816);
figure;
image(newT);
colormap(C/255);
axis image; 


figure;
plot(jsum_iter(:));

figure;
plot(jroc(:));

Can you tell me if there is a mistake on the third piece of code or is there something wrong in the last one?